compiler/plugins/input/StableHLO/Conversion/StableHLOToIREEInputDialects.cpp

// Copyright 2020 The IREE Authors
//
// Licensed 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

// Implements IREE-specific logic for lowering StableHLO dialect to
// IREE dialects: Linalg, Arith, Math, Tensor, Util, ML Program, etc.

#include "compiler/plugins/input/StableHLO/Conversion/LegalizeToLinalgUtils.h"
#include "compiler/plugins/input/StableHLO/Conversion/PassDetail.h"
#include "compiler/plugins/input/StableHLO/Conversion/Passes.h"
#include "compiler/plugins/input/StableHLO/Conversion/Preprocessing/Rewriters.h"
#include "compiler/plugins/input/StableHLO/Conversion/Rewriters.h"
#include "compiler/plugins/input/StableHLO/Conversion/TypeConversion.h"
#include "iree/compiler/Dialect/Flow/IR/FlowOps.h"
#include "iree/compiler/Dialect/LinalgExt/IR/LinalgExtDialect.h"
#include "iree/compiler/Dialect/Util/IR/UtilDialect.h"
#include "iree/compiler/Dialect/Util/IR/UtilOps.h"
#include "iree/compiler/Utils/ConversionUtils.h"
#include "llvm/ADT/SmallVector.h"
#include "mlir/Dialect/Arith/IR/Arith.h"
#include "mlir/Dialect/Complex/IR/Complex.h"
#include "mlir/Dialect/ControlFlow/IR/ControlFlowOps.h"
#include "mlir/Dialect/Func/IR/FuncOps.h"
#include "mlir/Dialect/Linalg/IR/Linalg.h"
#include "mlir/Dialect/Linalg/Transforms/Transforms.h"
#include "mlir/Dialect/MLProgram/IR/MLProgram.h"
#include "mlir/Dialect/Math/IR/Math.h"
#include "mlir/Dialect/MemRef/IR/MemRef.h"
#include "mlir/Dialect/Tensor/IR/Tensor.h"
#include "mlir/IR/BuiltinTypes.h"
#include "mlir/IR/MLIRContext.h"
#include "mlir/IR/PatternMatch.h"
#include "mlir/IR/TypeUtilities.h"
#include "mlir/Transforms/DialectConversion.h"
#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
#include "stablehlo/dialect/ChloOps.h"
#include "stablehlo/dialect/StablehloOps.h"

namespace mlir::iree_compiler::stablehlo {

#define GEN_PASS_DEF_CONVERTSTABLEHLOTOIREEINPUTDIALECTS
#include "compiler/plugins/input/StableHLO/Conversion/Passes.h.inc"

namespace {

/// Converts stablehlo.concatenate operation to extract_slice ops + insert_slice
/// ops.
struct ConcatenateOpConversion final
    : OpConversionPattern<mlir::stablehlo::ConcatenateOp> {
  using OpConversionPattern::OpConversionPattern;

  LogicalResult
  matchAndRewrite(mlir::stablehlo::ConcatenateOp op, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    auto resultType =
        getTypeConverter()->convertType<RankedTensorType>(op.getType());
    if (!resultType || !resultType.hasStaticShape()) {
      return rewriter.notifyMatchFailure(op,
                                         "expected static shape for output");
    }

    Location loc = op.getLoc();
    uint64_t dim = op.getDimension();
    int64_t rank = resultType.getRank();
    SmallVector<Value, 3> offsets;
    SmallVector<Value, 3> sizes;
    SmallVector<Value, 3> strides;

    for (int64_t i = 0; i < rank; ++i) {
      offsets.push_back(rewriter.create<arith::ConstantIndexOp>(loc, 0));
      sizes.push_back(rewriter.createOrFold<tensor::DimOp>(
          loc, adaptor.getOperands()[0], i));
      strides.push_back(rewriter.create<arith::ConstantIndexOp>(loc, 1));
    }

    Value resultDimSize = rewriter.create<arith::ConstantIndexOp>(loc, 0);
    for (Value arg : adaptor.getOperands()) {
      auto size = rewriter.createOrFold<tensor::DimOp>(loc, arg, dim);
      resultDimSize =
          rewriter.createOrFold<arith::AddIOp>(loc, resultDimSize, size);
    }
    sizes[dim] = resultDimSize;

    Value result = rewriter.create<tensor::EmptyOp>(
        loc, resultType.getShape(), resultType.getElementType());

    auto toOpFoldResult = [](Value v) -> OpFoldResult {
      auto op = v.getDefiningOp<arith::ConstantIndexOp>();
      if (!op)
        return v;
      return op.getValue();
    };

    Value accBound = rewriter.create<arith::ConstantIndexOp>(loc, 0);
    for (Value arg : adaptor.getOperands()) {
      offsets[dim] = accBound;
      sizes[dim] = rewriter.createOrFold<tensor::DimOp>(loc, arg, dim);
      result = rewriter.create<tensor::InsertSliceOp>(
          loc, arg, result, llvm::map_to_vector(offsets, toOpFoldResult),
          llvm::map_to_vector(sizes, toOpFoldResult),
          llvm::map_to_vector(strides, toOpFoldResult));
      accBound = rewriter.create<arith::AddIOp>(loc, accBound, sizes[dim]);
    }
    rewriter.replaceOp(op, result);
    return success();
  }
};

/// Creates coefficients based on DFT definition, see
/// https://en.wikipedia.org/wiki/Discrete_Fourier_transform.
Value getDFTMatmulCoeff(OpBuilder b, Location loc, RankedTensorType matrixType,
                        bool isRealPart) {
  // scale = 2 * pi / N
  double scale = 2 * M_PI / matrixType.getDimSize(0);

  SmallVector<Attribute> values;
  assert(matrixType.getRank() == 2 && "expected 2D matrix");
  for (auto i : llvm::seq<unsigned>(0, matrixType.getDimSize(0))) {
    for (auto j : llvm::seq<unsigned>(0, matrixType.getDimSize(1))) {
      double v = scale * i * j;
      if (isRealPart) {
        v = cos(v);
      } else {
        v = -sin(v);
      }
      values.push_back(b.getF32FloatAttr(v));
    }
  }
  return b.create<arith::ConstantOp>(
      loc, matrixType, DenseFPElementsAttr::get(matrixType, values));
}

Value createLinalgMatmulOnTensors(OpBuilder b, Location loc,
                                  RankedTensorType resultType, Value lhs,
                                  Value rhs) {
  Value zero = b.create<arith::ConstantOp>(
      loc, b.getZeroAttr(resultType.getElementType()));
  Value emptyTensor = b.create<mlir::tensor::EmptyOp>(
      loc, resultType.getShape(), resultType.getElementType(),
      /*dyn_size=*/ValueRange{});
  Value zeroTensor =
      b.create<linalg::FillOp>(loc, zero, emptyTensor).getResult(0);

  switch (llvm::cast<RankedTensorType>(lhs.getType()).getRank()) {
  case 1:
    return b
        .create<linalg::VecmatOp>(loc, TypeRange{resultType},
                                  ValueRange{lhs, rhs}, ValueRange{zeroTensor})
        .getResult(0);
  case 2:
    return b
        .create<linalg::MatmulOp>(loc, TypeRange{resultType},
                                  ValueRange{lhs, rhs}, ValueRange{zeroTensor})
        .getResult(0);
  default:
    assert(false && "unhandled matmul type");
    return Value();
  }
}

/// Converts stablehlo.fft operation to Linalg ops.
struct FftOpConversion final : OpConversionPattern<mlir::stablehlo::FftOp> {
  using OpConversionPattern::OpConversionPattern;

  LogicalResult
  matchAndRewrite(mlir::stablehlo::FftOp op, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    if (op.getFftType() != mlir::stablehlo::FftType::RFFT) {
      return rewriter.notifyMatchFailure(op,
                                         "non RFFT types are supported yet");
    }

    auto inputType = dyn_cast<RankedTensorType>(adaptor.getOperand().getType());
    if (!inputType || !inputType.hasStaticShape() || inputType.getRank() > 2) {
      return rewriter.notifyMatchFailure(op, "only static 1D or 2D dft ops");
    }

    int64_t rank = inputType.getRank();
    int64_t n = inputType.getDimSize(rank - 1);
    int64_t fftLength = op.getFftLength().front() / 2 + 1;

    Location loc = op.getLoc();
    auto matrixType =
        RankedTensorType::get({n, fftLength}, inputType.getElementType());
    auto resultType = RankedTensorType::get(
        llvm::cast<RankedTensorType>(op.getType()).getShape(),
        inputType.getElementType());

    Value realMatrix =
        getDFTMatmulCoeff(rewriter, loc, matrixType, /*isRealPart=*/true);
    Value real = createLinalgMatmulOnTensors(rewriter, loc, resultType,
                                             adaptor.getOperand(), realMatrix);

    Value imagMatrix =
        getDFTMatmulCoeff(rewriter, loc, matrixType, /*isRealPart=*/false);
    Value imag = createLinalgMatmulOnTensors(rewriter, loc, resultType,
                                             adaptor.getOperand(), imagMatrix);

    // Pack the results back to mlir::stablehlo::ComplexOp.
    rewriter.replaceOpWithNewOp<mlir::stablehlo::ComplexOp>(op, op.getType(),
                                                            real, imag);
    return success();
  }
};

struct OptimizationBarrierOpConversion final
    : OpConversionPattern<mlir::stablehlo::OptimizationBarrierOp> {
  using OpConversionPattern::OpConversionPattern;

  LogicalResult
  matchAndRewrite(mlir::stablehlo::OptimizationBarrierOp op, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    SmallVector<Value> outputs;
    for (Value operand : adaptor.getOperands()) {
      outputs.push_back(
          rewriter
              .create<IREE::Util::OptimizationBarrierOp>(op.getLoc(), operand)
              .getResult(0));
    }
    rewriter.replaceOp(op, outputs);
    return success();
  }
};

// Returns true if all attributes in the given dictionary are valid for IREE
// input dialects.
static bool isValidFuncAttr(DictionaryAttr attrs) {
  // TODO: switch to using a dialect-based exclusion list or some other way that
  // is not a big string table.
  for (auto attr : attrs) {
    if (attr.getName() == "tf.aliasing_output")
      return false;
  }
  return true;
}

// Adds iree.abi.encoding attributes for arguments and results when they have
// had their type changed during conversion.
static void setFuncEncodings(func::FuncOp funcOp, FunctionType oldFuncType,
                             FunctionType newFuncType) {
  auto encodingName = StringAttr::get(funcOp.getContext(), "iree.abi.encoding");
  for (auto [i, oldType, newType] :
       llvm::enumerate(oldFuncType.getInputs(), newFuncType.getInputs())) {
    if (oldType != newType)
      funcOp.setArgAttr(i, encodingName, TypeAttr::get(oldType));
  }
  for (auto [i, oldType, newType] :
       llvm::enumerate(oldFuncType.getResults(), newFuncType.getResults())) {
    if (oldType != newType)
      funcOp.setResultAttr(i, encodingName, TypeAttr::get(oldType));
  }
}

// Rewrites attributes on the function from ones coming from HLO-based frontends
// to the IREE supported versions.
static void rewriteFuncAttrs(func::FuncOp funcOp) {
  auto *context = funcOp.getContext();
  auto indexType = IndexType::get(context);
  auto abiOutputName = StringAttr::get(context, "iree.abi.output");
  auto aliasingOutputName = StringAttr::get(context, "tf.aliasing_output");
  auto rewriteAttrs = [&](DictionaryAttr &allAttrs) {
    SmallVector<NamedAttribute> newAttrs;
    newAttrs.reserve(allAttrs.size());
    for (auto attr : allAttrs) {
      if (attr.getName() == aliasingOutputName) {
        newAttrs.push_back({
            abiOutputName,
            IntegerAttr::get(indexType,
                             llvm::cast<IntegerAttr>(attr.getValue()).getInt()),
        });
      } else {
        newAttrs.push_back(attr);
      }
    }
    allAttrs = DictionaryAttr::get(context, newAttrs);
  };
  SmallVector<DictionaryAttr> argAttrs;
  funcOp.getAllArgAttrs(argAttrs);
  llvm::for_each(argAttrs, rewriteAttrs);
  funcOp.setAllArgAttrs(argAttrs);
  SmallVector<DictionaryAttr> resultAttrs;
  funcOp.getAllResultAttrs(resultAttrs);
  llvm::for_each(resultAttrs, rewriteAttrs);
  funcOp.setAllResultAttrs(resultAttrs);
}

// We need to convert func ops in order to convert types.
struct BuiltinFuncOpPattern final : OpConversionPattern<func::FuncOp> {
  using OpConversionPattern::OpConversionPattern;

  LogicalResult
  matchAndRewrite(func::FuncOp srcOp, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    FunctionType srcFuncType = srcOp.getFunctionType();
    TypeConverter::SignatureConversion signatureConversion(
        srcOp.getNumArguments());

    // Convert function arguments.
    for (auto [idx, inputTy] : llvm::enumerate(srcFuncType.getInputs())) {
      if (failed(getTypeConverter()->convertSignatureArg(
              idx, inputTy, signatureConversion))) {
        return rewriter.notifyMatchFailure(srcOp, "argument failed to convert");
      }
    }

    // Convert function results.
    SmallVector<Type> convertedResultTypes;
    if (failed(getTypeConverter()->convertTypes(srcFuncType.getResults(),
                                                convertedResultTypes))) {
      return rewriter.notifyMatchFailure(srcOp, "results failed to convert");
    }

    // Create new function with converted argument and result types.
    auto oldFuncType = srcOp.getFunctionType();
    auto newFuncType = mlir::FunctionType::get(
        srcOp.getContext(), signatureConversion.getConvertedTypes(),
        convertedResultTypes);

    // Update the function in place.
    rewriter.startOpModification(srcOp);
    srcOp.setType(newFuncType);
    rewriteFuncAttrs(srcOp);
    setFuncEncodings(srcOp, oldFuncType, newFuncType);

    // Tell the rewriter to convert the region signature.
    const TypeConverter &typeConverter = *getTypeConverter();
    if (failed(rewriter.convertRegionTypes(&srcOp.getBody(), typeConverter,
                                           &signatureConversion))) {
      return failure();
    }

    rewriter.finalizeOpModification(srcOp);
    return success();
  }
};

struct TensorEmptyPattern final : OpConversionPattern<tensor::EmptyOp> {
  using OpConversionPattern<tensor::EmptyOp>::OpConversionPattern;

  LogicalResult
  matchAndRewrite(tensor::EmptyOp op, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    auto oldType = cast<ShapedType>(op.getType());
    auto newType = getTypeConverter()->convertType(oldType);
    if (newType == oldType)
      return failure();

    if (!newType)
      return rewriter.notifyMatchFailure(op, "result type conversion failed");

    rewriter.replaceOpWithNewOp<tensor::EmptyOp>(
        op, oldType.getShape(),
        getTypeConverter()->convertType(oldType.getElementType()),
        op.getDynamicSizes());
    return success();
  }
};

struct GlobalOpPattern final : OpConversionPattern<ml_program::GlobalOp> {
  using OpConversionPattern<ml_program::GlobalOp>::OpConversionPattern;

  LogicalResult
  matchAndRewrite(ml_program::GlobalOp globalOp, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    Type oldType = globalOp.getType();
    Type newType = getTypeConverter()->convertType(oldType);
    if (newType == oldType)
      return failure();
    if (!newType) {
      return rewriter.notifyMatchFailure(globalOp,
                                         "result type conversion failed");
    }
    rewriter.modifyOpInPlace(globalOp, [&]() {
      globalOp.setType(newType);
      if (Attribute oldValue = globalOp.getValueAttr()) {
        globalOp.setValueAttr(
            convertAttribute(globalOp.getLoc(), oldValue, *getTypeConverter()));
      }
    });
    return success();
  }
};

template <typename T>
struct GenericTypeConvert final : ConversionPattern {
  GenericTypeConvert(StringRef rootName, TypeConverter &converter,
                     MLIRContext *context, PatternBenefit benefit = 0)
      : ConversionPattern(converter, rootName, benefit, context) {}

  GenericTypeConvert(TypeConverter &converter, MLIRContext *context,
                     PatternBenefit benefit = 0)
      : ConversionPattern(converter, T::getOperationName(), benefit, context) {}

  LogicalResult
  matchAndRewrite(Operation *op, ArrayRef<Value> operands,
                  ConversionPatternRewriter &rewriter) const override {
    llvm::SmallVector<NamedAttribute> newAttr;
    llvm::append_range(newAttr, op->getAttrs());

    llvm::SmallVector<Type> newResults;
    if (failed(getTypeConverter()->convertTypes(op->getResultTypes(),
                                                newResults))) {
      return rewriter.notifyMatchFailure(op, "result type conversion failed");
    }

    OperationState state(op->getLoc(), op->getName().getStringRef(), operands,
                         newResults, newAttr, op->getSuccessors());
    for (Region &r : op->getRegions()) {
      Region *newRegion = state.addRegion();
      rewriter.inlineRegionBefore(r, *newRegion, newRegion->begin());
      TypeConverter::SignatureConversion result(newRegion->getNumArguments());
      if (failed(getTypeConverter()->convertSignatureArgs(
              newRegion->getArgumentTypes(), result))) {
        return rewriter.notifyMatchFailure(op,
                                           "argument type conversion failed");
      }
      rewriter.applySignatureConversion(&newRegion->front(), result);
    }
    Operation *newOp = rewriter.create(state);
    rewriter.replaceOp(op, newOp->getResults());
    return success();
  }
};

Value scalarToTensor(OpBuilder &builder, Type /*type*/, ValueRange inputs,
                     Location loc) {
  assert(inputs.size() == 1);
  if (isa<ShapedType>(inputs.front().getType())) {
    return Value();
  }
  return builder
      .create<tensor::FromElementsOp>(
          loc, RankedTensorType::get({}, inputs.front().getType()),
          inputs.front())
      .getResult();
}

// Strips attributes from common StableHLO frontends (JAX, TF, etc) that are not
// used after conversion into the IREE input dialects. Leaving these attributes
// is confusing as they can become inconsistent during subsequent conversions or
// leak frontend details lower into the pipeline than should be allowed.
static void stripFrontendAttrs(mlir::ModuleOp moduleOp) {
  auto isAttrFiltered = [](NamedAttribute attr) {
    auto fullName = attr.getName().getValue();
    return fullName.starts_with("mhlo.") || fullName.starts_with("jax.") ||
           fullName.starts_with("tf.");
  };
  auto filterOpAttrs = [&](Operation *op) {
    SmallVector<NamedAttribute> newAttrs;
    for (auto attr : op->getDialectAttrs()) {
      if (!isAttrFiltered(attr))
        newAttrs.push_back(attr);
    }
    op->setDialectAttrs(newAttrs);
  };
  auto filterAttrDicts = [&](ArrayAttr allOldAttrs,
                             SmallVectorImpl<DictionaryAttr> &newAttrs) {
    if (!allOldAttrs)
      return false;
    for (auto oldAttrs : allOldAttrs.getAsRange<DictionaryAttr>()) {
      SmallVector<NamedAttribute> preservedAttrs;
      preservedAttrs.reserve(oldAttrs.size());
      for (auto attr : oldAttrs) {
        if (!isAttrFiltered(attr))
          preservedAttrs.push_back(attr);
      }
      newAttrs.push_back(
          DictionaryAttr::get(allOldAttrs.getContext(), preservedAttrs));
    }
    return true;
  };
  filterOpAttrs(moduleOp);
  for (auto callableOp : moduleOp.getOps<mlir::CallableOpInterface>()) {
    filterOpAttrs(callableOp);
    if (auto funcOp = dyn_cast<func::FuncOp>(callableOp.getOperation())) {
      SmallVector<DictionaryAttr> newArgAttrs;
      if (filterAttrDicts(funcOp.getAllArgAttrs(), newArgAttrs)) {
        funcOp.setAllArgAttrs(newArgAttrs);
      }
      SmallVector<DictionaryAttr> newResultAttrs;
      if (filterAttrDicts(funcOp.getAllResultAttrs(), newResultAttrs)) {
        funcOp.setAllResultAttrs(newResultAttrs);
      }
    }
  }
}

struct ConvertStableHloToIreeInputDialects final
    : impl::ConvertStableHloToIreeInputDialectsBase<
          ConvertStableHloToIreeInputDialects> {
  void getDependentDialects(DialectRegistry &registry) const override {
    registry.insert<
        IREE::Flow::FlowDialect, IREE::Util::UtilDialect, linalg::LinalgDialect,
        arith::ArithDialect, tensor::TensorDialect, shape::ShapeDialect,
        math::MathDialect, memref::MemRefDialect, complex::ComplexDialect>();
  }

  void runOnOperation() override {
    MLIRContext *context = &getContext();
    RewritePatternSet patterns(context);

    std::unique_ptr<TypeConverter> typeConverter =
        createStableHloToLinalgTypeConverter();
    typeConverter->addArgumentMaterialization(scalarToTensor);

    // Run stablehlo canonicalization patterns with a high benefit to avoid some
    // expensive expansions.
    populateCanonicalizationPatterns(context, &patterns, /*benefit=*/1024);

    populateStableHloToLinalgOnTensorsConversionPatterns(
        context, *typeConverter, &patterns);
    populateStableHloCollectivesConversionPatterns(context, *typeConverter,
                                                   &patterns);

    // TODO(#12678): Handle remaining complex ops.

    // TODO(*): expose patterns that do this much better from
    // iree/compiler/Dialect/Util/Transforms/ConvertPrimitiveType.cpp

    // Structural patterns (functions, cfg, terminators).
    patterns.add<BuiltinFuncOpPattern>(*typeConverter, context);
    patterns.add<GlobalOpPattern, TensorEmptyPattern>(*typeConverter, context);

    patterns.add<
        GenericTypeConvert<cf::CondBranchOp>, GenericTypeConvert<cf::BranchOp>,
        GenericTypeConvert<func::ReturnOp>, GenericTypeConvert<func::ReturnOp>,
        GenericTypeConvert<func::CallOp>,
        GenericTypeConvert<ml_program::GlobalLoadOp>,
        GenericTypeConvert<ml_program::GlobalLoadConstOp>,
        GenericTypeConvert<ml_program::GlobalStoreOp>,
        GenericTypeConvert<scf::ForOp>, GenericTypeConvert<scf::IfOp>,
        GenericTypeConvert<scf::YieldOp>, GenericTypeConvert<scf::ConditionOp>,
        GenericTypeConvert<scf::WhileOp>,
        GenericTypeConvert<tensor::FromElementsOp>,
        GenericTypeConvert<tensor::CollapseShapeOp>,
        GenericTypeConvert<tensor::ExpandShapeOp>,
        GenericTypeConvert<arith::IndexCastUIOp>,
        GenericTypeConvert<arith::SelectOp>>(*typeConverter, context);

    ConversionTarget target(*context);
    auto isIllegalType = [&](Type t) { return !typeConverter->isLegal(t); };
    auto isLegallyTypedOp = [&](Operation *op) -> bool {
      for (Type type : op->getResultTypes()) {
        if (isIllegalType(type))
          return false;
      }
      for (Type type : op->getOperandTypes()) {
        if (isIllegalType(type))
          return false;
      }
      return true;
    };

    target.addIllegalDialect<mlir::chlo::ChloDialect>();
    target.addIllegalDialect<mlir::stablehlo::StablehloDialect>();

    // Functions must have legal types.
    target.addDynamicallyLegalOp<func::FuncOp>([&](func::FuncOp funcOp) {
      if (auto attrs = funcOp.getAllArgAttrs()) {
        if (!llvm::all_of(attrs.getAsRange<DictionaryAttr>(),
                          isValidFuncAttr)) {
          return false;
        }
      }
      if (auto attrs = funcOp.getAllResultAttrs()) {
        if (!llvm::all_of(attrs.getAsRange<DictionaryAttr>(),
                          isValidFuncAttr)) {
          return false;
        }
      }
      for (Type type : funcOp.getFunctionType().getInputs()) {
        if (isIllegalType(type))
          return false;
      }
      for (Type type : funcOp.getFunctionType().getResults()) {
        if (isIllegalType(type))
          return false;
      }
      for (Block &block : funcOp.getFunctionBody()) {
        for (Type type : block.getArgumentTypes()) {
          if (isIllegalType(type))
            return false;
        }
      }
      return true;
    });
    target.addDynamicallyLegalOp<ml_program::GlobalOp>(
        [&](ml_program::GlobalOp op) {
          return typeConverter->isLegal(op.getType());
        });

    target.addDynamicallyLegalOp<tensor::EmptyOp>([&](tensor::EmptyOp op) {
      return typeConverter->isLegal(op.getType());
    });

    // Let the rest fall through.
    target.addLegalDialect<BuiltinDialect>();
    target.addLegalDialect<IREE::LinalgExt::IREELinalgExtDialect>();
    target.markUnknownOpDynamicallyLegal(isLegallyTypedOp);

    if (failed(applyPartialConversion(getOperation(), target,
                                      std::move(patterns)))) {
      return signalPassFailure();
    }

    {
      // Apply the patterns to remove unused operands and results.
      RewritePatternSet removeUnusedOperandsResultsPatterns(context);
      linalg::populateEraseUnusedOperandsAndResultsPatterns(
          removeUnusedOperandsResultsPatterns);
      if (failed(applyPatternsAndFoldGreedily(
              getOperation(),
              std::move(removeUnusedOperandsResultsPatterns)))) {
        return signalPassFailure();
      }
    }

    // Drop module/function attributes now that they are no longer required.
    stripFrontendAttrs(getOperation());
  }
};

} // namespace

void populateStableHloToLinalgOnTensorsConversionPatterns(
    MLIRContext *context, TypeConverter &typeConverter,
    RewritePatternSet *patterns) {
  // TODO(#5809): Drop ConcatenateOp lowering in favor of the upstream version
  //              then remove the PatternBenefit here
  patterns->add<ConcatenateOpConversion, FftOpConversion,
                OptimizationBarrierOpConversion>(typeConverter, context,
                                                 PatternBenefit{1000});

  populateStableHloToLinalgConversionPatterns(context, typeConverter, patterns,
                                              /*enablePrimitiveOps=*/false);
}

} // namespace mlir::iree_compiler::stablehlo