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"
 
#include <cassert>
 
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,
 
  /// Pass a pointer to the caller's object with no copy implied (byref).  The
  /// callee must not write through it.  Arguments only.
  IndirectAliased,
 
  /// 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 and IndirectAliased: 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;
 
  /// For Indirect and IndirectAliased: target address space of the pointer
  /// passed in place of the value.  Zero is the default address space.
  unsigned indirectAddrSpace = 0;
 
  /// For Direct with coercion: the byte offset within the original aggregate
  /// at which the coerced value lives.  Non-zero when the low eightbyte is
  /// NO_CLASS and the value is carried in a later eightbyte (x86-64 SysV).
  unsigned directOffset = 0;
 
  /// For Direct: the alignment of the argument's stack slot when the value
  /// is passed in memory, if the ABI sets one apart from the type's.
  llvm::MaybeAlign directAlign;
 
  /// How many integer and vector argument registers the value occupies.  Both
  /// zero means it travels in memory, which is also what a target whose
  /// classifier does not record the demand reports.
  unsigned neededIntRegs = 0;
  unsigned neededSseRegs = 0;
 
  /// Whether the value is passed as-is, so a rewriter can leave it alone.
  /// Only an uncoerced Direct with no stack-slot alignment qualifies.
  /// Extend counts as needing a rewrite even though it only adds an
  /// attribute, because the attribute changes observable behavior, and a
  /// Direct with a stack-slot alignment does too because it needs
  /// llvm.alignstack.
  bool isPassThrough() const {
    return kind == ArgKind::Direct && !coercedType && !directAlign;
  }
 
  /// 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 &&
           indirectAddrSpace == other.indirectAddrSpace &&
           directOffset == other.directOffset &&
           directAlign == other.directAlign &&
           neededIntRegs == other.neededIntRegs &&
           neededSseRegs == other.neededSseRegs;
  }
 
  static ArgClassification getDirect() {
    return getDirect(/*coerced=*/nullptr, /*offset=*/0);
  }
 
  static ArgClassification getDirect(Type coerced, unsigned offset,
                                     llvm::MaybeAlign align = std::nullopt) {
    // isPassThrough ignores directOffset, so an offset with no coerced
    // type to read at it would be silently ignored.
    assert((!offset || coerced) &&
           "a direct offset needs a coerced type to read at it");
    ArgClassification c;
    c.kind = ArgKind::Direct;
    c.coercedType = coerced;
    c.directOffset = offset;
    c.directAlign = align;
    return c;
  }
 
  static ArgClassification getIgnore() {
    ArgClassification c;
    c.kind = ArgKind::Ignore;
    return c;
  }
 
  static ArgClassification getIndirect(llvm::Align align, bool byVal,
                                       unsigned addrSpace) {
    ArgClassification c;
    c.kind = ArgKind::Indirect;
    c.indirectAlign = align;
    c.byVal = byVal;
    c.indirectAddrSpace = addrSpace;
    return c;
  }
 
  static ArgClassification getIndirectAliased(llvm::Align align,
                                              unsigned addrSpace) {
    ArgClassification c;
    c.kind = ArgKind::IndirectAliased;
    c.indirectAlign = align;
    c.indirectAddrSpace = addrSpace;
    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;
 
  /// Rewrite a single "fetch the next vararg" operation (e.g. C `va_arg`) to
  /// match how \p ac says the fetched type is passed at the ABI level.
  ///
  /// \p ac classifies only the one type being fetched, in isolation, with the
  /// whole register budget available.  A vararg fetch advances a runtime
  /// cursor (the platform va_list) through registers and then memory, so
  /// whether a given fetch lands in a register depends on how much of the
  /// budget earlier variadic arguments already consumed at run time, not on
  /// the fetch's static position.
  ///
  /// An implementation may erase \p vaArgOp and replace its result, so a
  /// caller walking the IR must collect the fetches before rewriting any of
  /// them.
  ///
  /// The default implementation reports failure, so a dialect that has not
  /// implemented vararg fetches does not need to override this.  An overrider
  /// that fails is responsible for emitting its own diagnostic.
  ///
  /// \param vaArgOp  The fetch to rewrite.  May be erased.
  /// \param ac       The ABI classification of the fetched type.
  /// \param builder  The OpBuilder to use for modifications.
  /// \returns success() if the operation was rewritten.
  virtual LogicalResult rewriteVAArg(Operation *vaArgOp,
                                     const ArgClassification &ac,
                                     OpBuilder &builder) {
    return failure();
  }
 
  /// Return the dialect namespace this context handles (e.g. "cir").
  virtual StringRef getDialectNamespace() const = 0;
};
 
} // namespace abi
} // namespace mlir
 
#endif // MLIR_ABI_ABIREWRITECONTEXT_H