ABIRewriteContext.h

Go to the documentation of this file.

//===- ABIRewriteContext.h - Dialect-specific ABI rewriting -----*- 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
//
//===----------------------------------------------------------------------===//
//
// This file defines ABIRewriteContext, the abstract interface for dialect-
// specific ABI lowering rewrites.  Each MLIR dialect that wants ABI lowering
// (CIR, FIR, etc.) provides a concrete subclass.
//
// ABIRewriteContext consumes ABI classification results and drives the
// creation of lowered function signatures, argument coercions, and call
// site rewrites using dialect-specific operations.
//
//===----------------------------------------------------------------------===//
 
#ifndef MLIR_ABI_ABIREWRITECONTEXT_H
#define MLIR_ABI_ABIREWRITECONTEXT_H
 
#include "mlir/IR/Builders.h"
#include "mlir/IR/Operation.h"
#include "mlir/IR/Types.h"
#include "mlir/IR/Value.h"
#include "mlir/Interfaces/FunctionInterfaces.h"
#include "llvm/ADT/STLExtras.h"
#include "llvm/Support/Alignment.h"
 
namespace mlir {
namespace abi {
 
/// Classification of how a single argument or return value should be
/// passed at the ABI level.
///
/// This is a dialect-agnostic representation.  It mirrors the kinds
/// found in the LLVM ABI library and in CIR's ABIArgInfo, but does
/// not depend on either.
enum class ArgKind : uint8_t {
  /// Pass directly in registers, possibly coerced to a different type.
  Direct,
 
  /// Like Direct, but with a sign/zero extension attribute.
  Extend,
 
  /// Pass indirectly via a pointer (sret for returns, byval for args).
  Indirect,
 
  /// Ignore (void return, empty struct).
  Ignore,
 
  /// Expand an aggregate into its constituent scalar fields.
  Expand,
};
 
/// Describes how a single argument or return value is passed after ABI
/// lowering.
struct ArgClassification {
  ArgKind kind = ArgKind::Direct;
 
  /// The ABI-coerced type, if different from the original.  Null means
  /// use the original type.
  Type coercedType = nullptr;
 
  /// For Indirect: alignment of the pointed-to object.
  llvm::Align indirectAlign = llvm::Align(1);
 
  /// For Extend: whether to sign-extend (true) or zero-extend (false).
  bool signExtend = false;
 
  /// For Direct: whether a struct coercion can be flattened into
  /// individual register-width arguments.
  bool canFlatten = true;
 
  /// For Indirect: whether the callee gets ownership (byval).
  bool byVal = false;
 
  /// Whether the value is passed as-is, so a rewriter can leave it alone.
  /// Only an uncoerced Direct qualifies.  Extend counts as needing a rewrite
  /// even though it only adds an attribute, because the attribute changes
  /// observable behavior.
  bool isPassThrough() const { return kind == ArgKind::Direct && !coercedType; }
 
  /// Whether two classifications describe the same wire format.  Every field
  /// participates, so a field added above must be added here as well.
  bool operator==(const ArgClassification &other) const {
    return kind == other.kind && coercedType == other.coercedType &&
           indirectAlign == other.indirectAlign &&
           signExtend == other.signExtend && canFlatten == other.canFlatten &&
           byVal == other.byVal;
  }
 
  static ArgClassification getDirect(Type coerced = nullptr) {
    ArgClassification c;
    c.kind = ArgKind::Direct;
    c.coercedType = coerced;
    return c;
  }
 
  static ArgClassification getIgnore() {
    ArgClassification c;
    c.kind = ArgKind::Ignore;
    return c;
  }
 
  static ArgClassification getIndirect(llvm::Align align, bool byVal = true) {
    ArgClassification c;
    c.kind = ArgKind::Indirect;
    c.indirectAlign = align;
    c.byVal = byVal;
    return c;
  }
 
  static ArgClassification getExtend(Type coerced, bool signExt) {
    ArgClassification c;
    c.kind = ArgKind::Extend;
    c.coercedType = coerced;
    c.signExtend = signExt;
    return c;
  }
};
 
/// Holds the full ABI classification for a function: return type and
/// all arguments.
struct FunctionClassification {
  ArgClassification returnInfo;
  SmallVector<ArgClassification> argInfos;
 
  /// Whether the classified return type was the source language's void.
  ///
  /// A void return classifies as Ignore, and so does a return the ABI drops,
  /// such as an empty record.  The two need opposite treatment: void is
  /// already its own wire form, while a dropped record return has to be
  /// rewritten to one.  returnInfo alone cannot tell them apart, so whoever
  /// produces the classification records it here, next to the return type it
  /// came from.  A consumer that re-derived it from something else could pair
  /// a classification with the wrong answer, and reading a dropped return as
  /// void means silently skipping the rewrite it needs.
  ///
  /// Left false when unknown, which costs a needless rewrite rather than a
  /// skipped one.
  bool returnsVoid = false;
 
  /// Whether any value in the signature is passed differently from how it is
  /// written, so a rewriter has work to do.
  bool needsRewrite() const {
    if (!returnsVoid && !returnInfo.isPassThrough())
      return true;
    return !llvm::all_of(argInfos, [](const ArgClassification &ac) {
      return ac.isPassThrough();
    });
  }
};
 
/// ABIRewriteContext is the abstract interface that each dialect
/// implements to perform ABI-specific rewrites on its operations.
///
/// The pass orchestrator calls these methods after ABI classification
/// to rewrite function definitions and call sites.
class ABIRewriteContext {
public:
  virtual ~ABIRewriteContext() = default;
 
  /// Rewrite a function definition to use ABI-lowered types.
  ///
  /// This creates a new function with the lowered signature, rewrites
  /// the function body to adapt between the ABI types and the
  /// original high-level types, and replaces the original function.
  ///
  /// \param funcOp  The function to rewrite (via FunctionOpInterface).
  /// \param fc      The ABI classification for this function.
  /// \param builder  The OpBuilder to use for modifications.
  /// \returns success() if the function was rewritten.
  virtual LogicalResult
  rewriteFunctionDefinition(FunctionOpInterface funcOp,
                            const FunctionClassification &fc,
                            OpBuilder &builder) = 0;
 
  /// Rewrite a call operation to match the callee's ABI-lowered
  /// signature.
  ///
  /// This coerces arguments, handles indirect returns (sret), and
  /// adapts the call result back to the original high-level type.
  ///
  /// \param callOp  The call operation to rewrite.
  /// \param fc      The ABI classification for the callee.
  /// \param builder  The OpBuilder to use for modifications.
  /// \returns success() if the call was rewritten.
  virtual LogicalResult rewriteCallSite(Operation *callOp,
                                        const FunctionClassification &fc,
                                        OpBuilder &builder) = 0;
 
  /// Return the dialect namespace this context handles (e.g. "cir").
  virtual StringRef getDialectNamespace() const = 0;
};
 
} // namespace abi
} // namespace mlir
 
#endif // MLIR_ABI_ABIREWRITECONTEXT_H