[mlir] Allow to specify target type in `convertType`

[kuhar]

Summary

Add a new helper function for the type converter that takes care of
casting to the target type.

This is to avoid bugs where an incorrect cast function is used after
type conversion, e.g., dyn_cast or cast. These are not guaranteed to
work when type conversion fails, or when type conversion succeeds but
the provided type converted returned a type that a conversion pattern
did not expect.

I saw this being an issue in some SPIR-V passes and in mlir-hlo.

Exercise the new function in a couple of passes. As a side-effect, this
also made the code more concise.

[mehdi_amini]

That makes the code quite nicer to read, thanks!

mlir/include/mlir/Transforms/DialectConversion.h

public:
   /// the type is removed and any usages of the existing value are expected to
   /// be removed during conversion.
   LogicalResult convertType(Type t, SmallVectorImpl<Type> &results);
 
   /// This hook simplifies defining 1-1 type conversions. This function returns
   /// the type to convert to on success, and a null type on failure.
   Type convertType(Type t);
 
+  /// Attempts a 1-1 type conversion, expecting the result type to be
+  /// `TargetType`. Returns the converted type cast to `TargetType` on success,
+  /// and a null type on conversion or cast failure.
+  template <typename TargetType>
+  TargetType convertType(Type t) {
+    return dyn_cast_or_null<TargetType>(convertType(t));
+  }
+
   /// Convert the given set of types, filling 'results' as necessary. This
   /// returns failure if the conversion of any of the types fails, success
   /// otherwise.
   LogicalResult convertTypes(TypeRange types, SmallVectorImpl<Type> &results);
 
   /// Return true if the given type is legal for this type converter, i.e. the
   /// type converts to itself.
   bool isLegal(Type type);

mlir/lib/Conversion/MemRefToSPIRV/MemRefToSPIRV.cpp

 //===- MemRefToSPIRV.cpp - MemRef to SPIR-V Patterns ----------------------===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 //
 // This file implements patterns to convert MemRef dialect to SPIR-V dialect.
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/Arith/IR/Arith.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/Dialect/SPIRV/IR/SPIRVEnums.h"
 #include "mlir/Dialect/SPIRV/IR/SPIRVOps.h"
+#include "mlir/Dialect/SPIRV/IR/SPIRVTypes.h"
 #include "mlir/Dialect/SPIRV/Transforms/SPIRVConversion.h"
 #include "mlir/IR/BuiltinTypes.h"
 #include "llvm/Support/Debug.h"
 #include <optional>
 
 #define DEBUG_TYPE "memref-to-spirv-pattern"
 
 using namespace mlir;
IntLoadOpPattern::matchAndRewrite(memref::LoadOp loadOp, OpAdaptor adaptor,
 
   if (!accessChain)
     return failure();
 
   int srcBits = memrefType.getElementType().getIntOrFloatBitWidth();
   bool isBool = srcBits == 1;
   if (isBool)
     srcBits = typeConverter.getOptions().boolNumBits;
-  Type pointeeType = typeConverter.convertType(memrefType)
-                         .cast<spirv::PointerType>()
-                         .getPointeeType();
+
+  auto pointerType = typeConverter.convertType<spirv::PointerType>(memrefType);
+  if (!pointerType)
+    return rewriter.notifyMatchFailure(loadOp, "failed to convert memref type");
+
+  Type pointeeType = pointerType.getPointeeType();
   Type dstType;
   if (typeConverter.allows(spirv::Capability::Kernel)) {
     if (auto arrayType = pointeeType.dyn_cast<spirv::ArrayType>())
       dstType = arrayType.getElementType();
     else
       dstType = pointeeType;
   } else {
     // For Vulkan we need to extract element from wrapping struct and array.
if (!accessChain)
     return failure();
 
   int srcBits = memrefType.getElementType().getIntOrFloatBitWidth();
 
   bool isBool = srcBits == 1;
   if (isBool)
     srcBits = typeConverter.getOptions().boolNumBits;
 
-  Type pointeeType = typeConverter.convertType(memrefType)
-                         .cast<spirv::PointerType>()
-                         .getPointeeType();
+  auto pointerType = typeConverter.convertType<spirv::PointerType>(memrefType);
+  if (!pointerType)
+    return rewriter.notifyMatchFailure(storeOp,
+                                       "failed to convert memref type");
+
+  Type pointeeType = pointerType.getPointeeType();
   Type dstType;
   if (typeConverter.allows(spirv::Capability::Kernel)) {
     if (auto arrayType = pointeeType.dyn_cast<spirv::ArrayType>())
       dstType = arrayType.getElementType();
     else
       dstType = pointeeType;
   } else {
     // For Vulkan we need to extract element from wrapping struct and array.

mlir/lib/Dialect/Arith/Transforms/EmulateWideInt.cpp

 //===- EmulateWideInt.cpp - Wide integer operation emulation ----*- C++ -*-===//
 //
 // 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/Transforms/Passes.h"
 
 #include "mlir/Dialect/Arith/IR/Arith.h"
 #include "mlir/Dialect/Arith/Transforms/WideIntEmulationConverter.h"
 #include "mlir/Dialect/Func/IR/FuncOps.h"
 #include "mlir/Dialect/Func/Transforms/FuncConversions.h"
 #include "mlir/Dialect/Vector/IR/VectorOps.h"
+#include "mlir/IR/BuiltinTypes.h"
 #include "mlir/IR/TypeUtilities.h"
 #include "mlir/Support/LogicalResult.h"
 #include "mlir/Transforms/DialectConversion.h"
 #include "llvm/ADT/APInt.h"
 #include "llvm/Support/FormatVariadic.h"
 #include "llvm/Support/MathExtras.h"
 #include <cassert>
 

 
 struct ConvertConstant final : OpConversionPattern<arith::ConstantOp> {
   using OpConversionPattern::OpConversionPattern;
 
   LogicalResult
   matchAndRewrite(arith::ConstantOp op, OpAdaptor,
                   ConversionPatternRewriter &rewriter) const override {
     Type oldType = op.getType();
-    auto newType = getTypeConverter()->convertType(oldType).cast<VectorType>();
+    auto newType = getTypeConverter()->convertType<VectorType>(oldType);
+    if (!newType)
+      return rewriter.notifyMatchFailure(
+          op, llvm::formatv("unsupported type: {0}", op.getType()));
+
     unsigned newBitWidth = newType.getElementTypeBitWidth();
     Attribute oldValue = op.getValueAttr();
 
     if (auto intAttr = oldValue.dyn_cast<IntegerAttr>()) {
       auto [low, high] = getHalves(intAttr.getValue(), newBitWidth);
       auto newAttr = DenseElementsAttr::get(newType, {low, high});
       rewriter.replaceOpWithNewOp<arith::ConstantOp>(op, newAttr);
       return success();

 
 struct ConvertAddI final : OpConversionPattern<arith::AddIOp> {
   using OpConversionPattern::OpConversionPattern;
 
   LogicalResult
   matchAndRewrite(arith::AddIOp op, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {
     Location loc = op->getLoc();
-    auto newTy = getTypeConverter()
-                     ->convertType(op.getType())
-                     .dyn_cast_or_null<VectorType>();
+    auto newTy = getTypeConverter()->convertType<VectorType>(op.getType());
     if (!newTy)
       return rewriter.notifyMatchFailure(
           loc, llvm::formatv("unsupported type: {0}", op.getType()));
 
     Type newElemTy = reduceInnermostDim(newTy);
 
     auto [lhsElem0, lhsElem1] =
         extractLastDimHalves(rewriter, loc, adaptor.getLhs());

 struct ConvertBitwiseBinary final : OpConversionPattern<BinaryOp> {
   using OpConversionPattern<BinaryOp>::OpConversionPattern;
   using OpAdaptor = typename OpConversionPattern<BinaryOp>::OpAdaptor;
 
   LogicalResult
   matchAndRewrite(BinaryOp op, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {
     Location loc = op->getLoc();
-    auto newTy = this->getTypeConverter()
-                     ->convertType(op.getType())
-                     .template dyn_cast_or_null<VectorType>();
+    auto newTy = this->getTypeConverter()->template convertType<VectorType>(
+        op.getType());
     if (!newTy)
       return rewriter.notifyMatchFailure(
           loc, llvm::formatv("unsupported type: {0}", op.getType()));
 
     auto [lhsElem0, lhsElem1] =
         extractLastDimHalves(rewriter, loc, adaptor.getLhs());
     auto [rhsElem0, rhsElem1] =
         extractLastDimHalves(rewriter, loc, adaptor.getRhs());

 
 struct ConvertCmpI final : OpConversionPattern<arith::CmpIOp> {
   using OpConversionPattern::OpConversionPattern;
 
   LogicalResult
   matchAndRewrite(arith::CmpIOp op, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {
     Location loc = op->getLoc();
-    auto inputTy = getTypeConverter()
-                       ->convertType(op.getLhs().getType())
-                       .dyn_cast_or_null<VectorType>();
+    auto inputTy =
+        getTypeConverter()->convertType<VectorType>(op.getLhs().getType());
     if (!inputTy)
       return rewriter.notifyMatchFailure(
           loc, llvm::formatv("unsupported type: {0}", op.getType()));
 
     arith::CmpIPredicate highPred = adaptor.getPredicate();
     arith::CmpIPredicate lowPred = toUnsignedPredicate(highPred);
 
     auto [lhsElem0, lhsElem1] =

 
 struct ConvertMulI final : OpConversionPattern<arith::MulIOp> {
   using OpConversionPattern::OpConversionPattern;
 
   LogicalResult
   matchAndRewrite(arith::MulIOp op, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {
     Location loc = op->getLoc();
-    auto newTy = getTypeConverter()
-                     ->convertType(op.getType())
-                     .dyn_cast_or_null<VectorType>();
+    auto newTy = getTypeConverter()->convertType<VectorType>(op.getType());
     if (!newTy)
       return rewriter.notifyMatchFailure(
           loc, llvm::formatv("unsupported type: {0}", op.getType()));
 
     auto [lhsElem0, lhsElem1] =
         extractLastDimHalves(rewriter, loc, adaptor.getLhs());
     auto [rhsElem0, rhsElem1] =
         extractLastDimHalves(rewriter, loc, adaptor.getRhs());

 
 struct ConvertExtSI final : OpConversionPattern<arith::ExtSIOp> {
   using OpConversionPattern::OpConversionPattern;
 
   LogicalResult
   matchAndRewrite(arith::ExtSIOp op, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {
     Location loc = op->getLoc();
-    auto newTy = getTypeConverter()
-                     ->convertType(op.getType())
-                     .dyn_cast_or_null<VectorType>();
+    auto newTy = getTypeConverter()->convertType<VectorType>(op.getType());
     if (!newTy)
       return rewriter.notifyMatchFailure(
           loc, llvm::formatv("unsupported type: {0}", op.getType()));
 
     Type newResultComponentTy = reduceInnermostDim(newTy);
 
     // Sign-extend the input value to determine the low half of the result.
     // Then, check if the low half is negative, and sign-extend the comparison

 
 struct ConvertExtUI final : OpConversionPattern<arith::ExtUIOp> {
   using OpConversionPattern::OpConversionPattern;
 
   LogicalResult
   matchAndRewrite(arith::ExtUIOp op, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {
     Location loc = op->getLoc();
-    auto newTy = getTypeConverter()
-                     ->convertType(op.getType())
-                     .dyn_cast_or_null<VectorType>();
+    auto newTy = getTypeConverter()->convertType<VectorType>(op.getType());
     if (!newTy)
       return rewriter.notifyMatchFailure(
           loc, llvm::formatv("unsupported type: {0}", op.getType()));
 
     Type newResultComponentTy = reduceInnermostDim(newTy);
 
     // Zero-extend the input value to determine the low half of the result.
     // The high half is always zero.
struct ConvertIndexCastIntToIndex final : OpConversionPattern<CastOp> {
   matchAndRewrite(CastOp op, typename CastOp::Adaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {
     Type resultType = op.getType();
     if (!isIndexOrIndexVector(resultType))
       return failure();
 
     Location loc = op.getLoc();
     Type inType = op.getIn().getType();
-    auto newInTy = this->getTypeConverter()
-                       ->convertType(inType)
-                       .template dyn_cast_or_null<VectorType>();
+    auto newInTy =
+        this->getTypeConverter()->template convertType<VectorType>(inType);
     if (!newInTy)
       return rewriter.notifyMatchFailure(
           loc, llvm::formatv("unsupported type: {0}", inType));
 
     // Discard the high half of the input truncating the original value.
     Value extracted = extractLastDimSlice(rewriter, loc, adaptor.getIn(), 0);
     extracted = dropTrailingX1Dim(rewriter, loc, extracted);
     rewriter.replaceOpWithNewOp<CastOp>(op, resultType, extracted);
matchAndRewrite(CastOp op, typename CastOp::Adaptor adaptor,
     if (!isIndexOrIndexVector(inType))
       return failure();
 
     Location loc = op.getLoc();
     auto *typeConverter =
         this->template getTypeConverter<arith::WideIntEmulationConverter>();
 
     Type resultType = op.getType();
-    auto newTy = typeConverter->convertType(resultType)
-                     .template dyn_cast_or_null<VectorType>();
+    auto newTy = typeConverter->template convertType<VectorType>(resultType);
     if (!newTy)
       return rewriter.notifyMatchFailure(
           loc, llvm::formatv("unsupported type: {0}", resultType));
 
     // Emit an index cast over the matching narrow type.
     Type narrowTy =
         rewriter.getIntegerType(typeConverter->getMaxTargetIntBitWidth());
     if (auto vecTy = resultType.dyn_cast<VectorType>())

 
 struct ConvertSelect final : OpConversionPattern<arith::SelectOp> {
   using OpConversionPattern::OpConversionPattern;
 
   LogicalResult
   matchAndRewrite(arith::SelectOp op, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {
     Location loc = op->getLoc();
-    auto newTy = getTypeConverter()
-                     ->convertType(op.getType())
-                     .dyn_cast_or_null<VectorType>();
+    auto newTy = getTypeConverter()->convertType<VectorType>(op.getType());
     if (!newTy)
       return rewriter.notifyMatchFailure(
           loc, llvm::formatv("unsupported type: {0}", op.getType()));
 
     auto [trueElem0, trueElem1] =
         extractLastDimHalves(rewriter, loc, adaptor.getTrueValue());
     auto [falseElem0, falseElem1] =
         extractLastDimHalves(rewriter, loc, adaptor.getFalseValue());
struct ConvertShLI final : OpConversionPattern<arith::ShLIOp> {
   using OpConversionPattern::OpConversionPattern;
 
   LogicalResult
   matchAndRewrite(arith::ShLIOp op, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {
     Location loc = op->getLoc();
 
     Type oldTy = op.getType();
-    auto newTy =
-        getTypeConverter()->convertType(oldTy).dyn_cast_or_null<VectorType>();
+    auto newTy = getTypeConverter()->convertType<VectorType>(oldTy);
     if (!newTy)
       return rewriter.notifyMatchFailure(
           loc, llvm::formatv("unsupported type: {0}", op.getType()));
 
     Type newOperandTy = reduceInnermostDim(newTy);
     // `oldBitWidth` == `2 * newBitWidth`
     unsigned newBitWidth = newTy.getElementTypeBitWidth();
 
struct ConvertShRUI final : OpConversionPattern<arith::ShRUIOp> {
   using OpConversionPattern::OpConversionPattern;
 
   LogicalResult
   matchAndRewrite(arith::ShRUIOp op, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {
     Location loc = op->getLoc();
 
     Type oldTy = op.getType();
-    auto newTy =
-        getTypeConverter()->convertType(oldTy).dyn_cast_or_null<VectorType>();
+    auto newTy = getTypeConverter()->convertType<VectorType>(oldTy);
     if (!newTy)
       return rewriter.notifyMatchFailure(
           loc, llvm::formatv("unsupported type: {0}", op.getType()));
 
     Type newOperandTy = reduceInnermostDim(newTy);
     // `oldBitWidth` == `2 * newBitWidth`
     unsigned newBitWidth = newTy.getElementTypeBitWidth();
 
struct ConvertShRSI final : OpConversionPattern<arith::ShRSIOp> {
   using OpConversionPattern::OpConversionPattern;
 
   LogicalResult
   matchAndRewrite(arith::ShRSIOp op, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {
     Location loc = op->getLoc();
 
     Type oldTy = op.getType();
-    auto newTy =
-        getTypeConverter()->convertType(oldTy).dyn_cast_or_null<VectorType>();
+    auto newTy = getTypeConverter()->convertType<VectorType>(oldTy);
     if (!newTy)
       return rewriter.notifyMatchFailure(
           loc, llvm::formatv("unsupported type: {0}", op.getType()));
 
     Value lhsElem1 = extractLastDimSlice(rewriter, loc, adaptor.getLhs(), 1);
     Value rhsElem0 = extractLastDimSlice(rewriter, loc, adaptor.getRhs(), 0);
 
     Type narrowTy = rhsElem0.getType();
struct ConvertSIToFP final : OpConversionPattern<arith::SIToFPOp> {
 
   LogicalResult
   matchAndRewrite(arith::SIToFPOp op, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {
     Location loc = op.getLoc();
 
     Value in = op.getIn();
     Type oldTy = in.getType();
-    auto newTy =
-        dyn_cast_or_null<VectorType>(getTypeConverter()->convertType(oldTy));
+    auto newTy = getTypeConverter()->convertType<VectorType>(oldTy);
     if (!newTy)
       return rewriter.notifyMatchFailure(
           loc, llvm::formatv("unsupported type: {0}", oldTy));
 
     unsigned oldBitWidth = getElementTypeOrSelf(oldTy).getIntOrFloatBitWidth();
     Value zeroCst = createScalarOrSplatConstant(rewriter, loc, oldTy, 0);
     Value oneCst = createScalarOrSplatConstant(rewriter, loc, oldTy, 1);
     Value allOnesCst = createScalarOrSplatConstant(
struct ConvertUIToFP final : OpConversionPattern<arith::UIToFPOp> {
   using OpConversionPattern::OpConversionPattern;
 
   LogicalResult
   matchAndRewrite(arith::UIToFPOp op, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {
     Location loc = op.getLoc();
 
     Type oldTy = op.getIn().getType();
-    auto newTy =
-        dyn_cast_or_null<VectorType>(getTypeConverter()->convertType(oldTy));
+    auto newTy = getTypeConverter()->convertType<VectorType>(oldTy);
     if (!newTy)
       return rewriter.notifyMatchFailure(
           loc, llvm::formatv("unsupported type: {0}", oldTy));
     unsigned newBitWidth = newTy.getElementTypeBitWidth();
 
     auto [low, hi] = extractLastDimHalves(rewriter, loc, adaptor.getIn());
     Value lowInt = dropTrailingX1Dim(rewriter, loc, low);
     Value hiInt = dropTrailingX1Dim(rewriter, loc, hi);