[CIR] Add support for __builtin_assume_aligned

Author: Lancern

clang/include/clang/CIR/Dialect/IR/CIROps.td

@@ -3143,6 +3143,45 @@ def CIR_AssumeOp : CIR_Op<"assume"> {
   }];
 }
 
+def CIR_AssumeAlignedOp : CIR_Op<"assume_aligned", [
+  Pure, AllTypesMatch<["pointer", "result"]>
+]> {
+  let summary = "Tell the optimizer that a pointer is aligned";
+  let description = [{
+    The `cir.assume_aligned` operation takes two or three arguments. The first
+    argument `pointer` gives the pointer value whose alignment is to be assumed,
+    and the second argument `align` is an integer attribute that gives the
+    assumed alignment.
+
+    The `offset` argument is optional. If given, it represents misalignment
+    offset. When it's present, this operation tells the optimizer that the
+    pointer is always misaligned to the alignment by `offset` bytes, a.k.a. the
+    pointer yielded by `(char *)pointer - offset` is aligned to the specified
+    alignment. Note that the `offset` argument is an SSA value rather than an
+    attribute, which means that you could pass a dynamically determined value
+    as the mialignment offset.
+
+    The result of this operation has the same value as the `pointer` argument,
+    but it additionally carries any alignment information indicated by this
+    operation.
+
+    This operation corresponds to the `__builtin_assume_aligned` builtin
+    function.
+  }];
+
+  let arguments = (ins CIR_PointerType:$pointer,
+                       I64Attr:$alignment,
+                       Optional<CIR_IntType>:$offset);
+  let results = (outs CIR_PointerType:$result);
+
+  let assemblyFormat = [{
+    $pointer
+    `alignment` $alignment
+    (`[` `offset` $offset^ `:` type($offset) `]`)?
+    `:` qualified(type($pointer)) attr-dict
+  }];
+}
+
 def CIR_AssumeSepStorageOp : CIR_Op<"assume_separate_storage", [
   SameTypeOperands
 ]> {

clang/lib/CIR/CodeGen/CIRGenBuiltin.cpp

@@ -129,6 +129,23 @@ RValue CIRGenFunction::emitBuiltinExpr(const GlobalDecl &gd, unsigned builtinID,
     return RValue::get(nullptr);
   }
 
+  case Builtin::BI__builtin_assume_aligned: {
+    const Expr *ptrExpr = e->getArg(0);
+    mlir::Value ptrValue = emitScalarExpr(ptrExpr);
+    mlir::Value offsetValue =
+        (e->getNumArgs() > 2) ? emitScalarExpr(e->getArg(2)) : nullptr;
+
+    const Expr *alignmentExpr = e->getArg(1);
+    mlir::Attribute alignmentAttr = ConstantEmitter(*this).emitAbstract(
+        alignmentExpr, alignmentExpr->getType());
+    int64_t alignment = mlir::cast<cir::IntAttr>(alignmentAttr).getSInt();
+
+    mlir::Value result = emitAlignmentAssumption(
+        ptrValue, ptrExpr, ptrExpr->getExprLoc(),
+        builder.getI64IntegerAttr(alignment), offsetValue);
+    return RValue::get(result);
+  }
+
   case Builtin::BI__builtin_complex: {
     mlir::Value real = emitScalarExpr(e->getArg(0));
     mlir::Value imag = emitScalarExpr(e->getArg(1));

clang/lib/CIR/CodeGen/CIRGenConstantEmitter.h

@@ -90,6 +90,7 @@ class ConstantEmitter {
   /// asserting that it succeeded.  This is only safe to do when the
   /// expression is known to be a constant expression with either a fairly
   /// simple type or a known simple form.
+  mlir::Attribute emitAbstract(const Expr *e, QualType ty);
   mlir::Attribute emitAbstract(SourceLocation loc, const APValue &value,
                                QualType t);
 
@@ -100,6 +101,7 @@ class ConstantEmitter {
   // functions and classes.
 
   mlir::Attribute tryEmitPrivateForVarInit(const VarDecl &d);
+  mlir::TypedAttr tryEmitPrivate(const Expr *e, QualType ty);
 
   mlir::Attribute tryEmitPrivate(const APValue &value, QualType destType);
   mlir::Attribute tryEmitPrivateForMemory(const APValue &value, QualType t);

clang/lib/CIR/CodeGen/CIRGenExprConstant.cpp

@@ -650,6 +650,38 @@ mlir::Attribute ConstantEmitter::tryEmitPrivateForMemory(const APValue &value,
   return (c ? emitForMemory(c, destType) : nullptr);
 }
 
+mlir::TypedAttr ConstantEmitter::tryEmitPrivate(const Expr *e, QualType ty) {
+  assert(!ty->isVoidType() && "can't emit a void constant");
+
+  if (mlir::Attribute attr =
+          ConstExprEmitter(*this).Visit(const_cast<Expr *>(e), ty))
+    return mlir::cast<mlir::TypedAttr>(attr);
+
+  Expr::EvalResult result;
+
+  bool success = false;
+  if (ty->isReferenceType())
+    success = e->EvaluateAsLValue(result, cgm.getASTContext());
+  else
+    success =
+        e->EvaluateAsRValue(result, cgm.getASTContext(), inConstantContext);
+
+  if (success && !result.hasSideEffects()) {
+    mlir::Attribute attr = tryEmitPrivate(result.Val, ty);
+    return mlir::cast<mlir::TypedAttr>(attr);
+  }
+
+  return nullptr;
+}
+
+mlir::Attribute ConstantEmitter::emitAbstract(const Expr *e, QualType ty) {
+  AbstractStateRAII state(*this, true);
+  auto attr = mlir::cast<mlir::Attribute>(tryEmitPrivate(e, ty));
+  if (!attr)
+    cgm.errorNYI(e->getExprLoc(), "emitAbstract failed, emit null constant");
+  return attr;
+}
+
 mlir::Attribute ConstantEmitter::emitAbstract(SourceLocation loc,
                                               const APValue &value,
                                               QualType destType) {

clang/lib/CIR/CodeGen/CIRGenFunction.cpp

@@ -931,6 +931,24 @@ CIRGenFunction::emitArrayLength(const clang::ArrayType *origArrayType,
   return builder.getConstInt(*currSrcLoc, SizeTy, countFromCLAs);
 }
 
+mlir::Value CIRGenFunction::emitAlignmentAssumption(
+    mlir::Value ptrValue, QualType ty, SourceLocation loc,
+    SourceLocation assumptionLoc, mlir::IntegerAttr alignment,
+    mlir::Value offsetValue) {
+  assert(!cir::MissingFeatures::sanitizers());
+  return cir::AssumeAlignedOp::create(builder, getLoc(assumptionLoc), ptrValue,
+                                      alignment, offsetValue);
+}
+
+mlir::Value CIRGenFunction::emitAlignmentAssumption(
+    mlir::Value ptrValue, const Expr *expr, SourceLocation assumptionLoc,
+    mlir::IntegerAttr alignment, mlir::Value offsetValue) {
+  QualType ty = expr->getType();
+  SourceLocation loc = expr->getExprLoc();
+  return emitAlignmentAssumption(ptrValue, ty, loc, assumptionLoc, alignment,
+                                 offsetValue);
+}
+
 // TODO(cir): Most of this function can be shared between CIRGen
 // and traditional LLVM codegen
 void CIRGenFunction::emitVariablyModifiedType(QualType type) {

clang/lib/CIR/CodeGen/CIRGenFunction.h

@@ -825,6 +825,18 @@ class CIRGenFunction : public CIRGenTypeCache {
   /// ----------------------
   /// CIR emit functions
   /// ----------------------
+public:
+  mlir::Value emitAlignmentAssumption(mlir::Value ptrValue, QualType ty,
+                                      SourceLocation loc,
+                                      SourceLocation assumptionLoc,
+                                      mlir::IntegerAttr alignment,
+                                      mlir::Value offsetValue = nullptr);
+
+  mlir::Value emitAlignmentAssumption(mlir::Value ptrValue, const Expr *expr,
+                                      SourceLocation assumptionLoc,
+                                      mlir::IntegerAttr alignment,
+                                      mlir::Value offsetValue = nullptr);
+
 private:
   void emitAndUpdateRetAlloca(clang::QualType type, mlir::Location loc,
                               clang::CharUnits alignment);

clang/lib/CIR/Lowering/DirectToLLVM/LowerToLLVM.cpp

@@ -460,6 +460,28 @@ mlir::LogicalResult CIRToLLVMAssumeOpLowering::matchAndRewrite(
   return mlir::success();
 }
 
+mlir::LogicalResult CIRToLLVMAssumeAlignedOpLowering::matchAndRewrite(
+    cir::AssumeAlignedOp op, OpAdaptor adaptor,
+    mlir::ConversionPatternRewriter &rewriter) const {
+  SmallVector<mlir::Value, 3> opBundleArgs{adaptor.getPointer()};
+
+  auto alignment = rewriter.create<mlir::LLVM::ConstantOp>(
+      op.getLoc(), rewriter.getI64Type(), op.getAlignment());
+  opBundleArgs.push_back(alignment);
+
+  if (mlir::Value offset = adaptor.getOffset())
+    opBundleArgs.push_back(offset);
+
+  auto cond = rewriter.create<mlir::LLVM::ConstantOp>(op.getLoc(),
+                                                      rewriter.getI1Type(), 1);
+  rewriter.create<mlir::LLVM::AssumeOp>(op.getLoc(), cond, "align",
+                                        opBundleArgs);
+  rewriter.replaceAllUsesWith(op, op.getPointer());
+  rewriter.eraseOp(op);
+
+  return mlir::success();
+}
+
 mlir::LogicalResult CIRToLLVMAssumeSepStorageOpLowering::matchAndRewrite(
     cir::AssumeSepStorageOp op, OpAdaptor adaptor,
     mlir::ConversionPatternRewriter &rewriter) const {
@@ -2168,6 +2190,7 @@ void ConvertCIRToLLVMPass::runOnOperation() {
   patterns.add<
       // clang-format off
                CIRToLLVMAssumeOpLowering,
+               CIRToLLVMAssumeAlignedOpLowering,
                CIRToLLVMAssumeSepStorageOpLowering,
                CIRToLLVMBaseClassAddrOpLowering,
                CIRToLLVMBinOpLowering,

clang/lib/CIR/Lowering/DirectToLLVM/LowerToLLVM.h

@@ -44,6 +44,16 @@ class CIRToLLVMAssumeOpLowering
                   mlir::ConversionPatternRewriter &) const override;
 };
 
+class CIRToLLVMAssumeAlignedOpLowering
+    : public mlir::OpConversionPattern<cir::AssumeAlignedOp> {
+public:
+  using mlir::OpConversionPattern<cir::AssumeAlignedOp>::OpConversionPattern;
+
+  mlir::LogicalResult
+  matchAndRewrite(cir::AssumeAlignedOp op, OpAdaptor,
+                  mlir::ConversionPatternRewriter &) const override;
+};
+
 class CIRToLLVMAssumeSepStorageOpLowering
     : public mlir::OpConversionPattern<cir::AssumeSepStorageOp> {
 public:

clang/test/CIR/CodeGen/builtin_call.cpp

@@ -111,6 +111,38 @@ void assume(bool arg) {
 // OGCG:   call void @llvm.assume(i1 %{{.+}})
 // OGCG: }
 
+void *assume_aligned(void *ptr) {
+  return __builtin_assume_aligned(ptr, 16);
+}
+
+// CIR: @_Z14assume_alignedPv
+// CIR:   %{{.+}} = cir.assume_aligned %{{.+}} alignment 16 : !cir.ptr<!void>
+// CIR: }
+
+// LLVM: @_Z14assume_alignedPv
+// LLVM:   call void @llvm.assume(i1 true) [ "align"(ptr %{{.+}}, i64 16) ]
+// LLVM: }
+
+// OGCG: @_Z14assume_alignedPv
+// OGCG:   call void @llvm.assume(i1 true) [ "align"(ptr %{{.+}}, i64 16) ]
+// OGCG: }
+
+void *assume_aligned_misalignment(void *ptr, unsigned misalignment) {
+  return __builtin_assume_aligned(ptr, 16, misalignment);
+}
+
+// CIR: @_Z27assume_aligned_misalignmentPvj
+// CIR:   %{{.+}} = cir.assume_aligned %{{.+}} alignment 16[offset %{{.+}} : !u64i] : !cir.ptr<!void>
+// CIR: }
+
+// LLVM: @_Z27assume_aligned_misalignmentPvj
+// LLVM:   call void @llvm.assume(i1 true) [ "align"(ptr %{{.+}}, i64 16, i64 %{{.+}}) ]
+// LLVM: }
+
+// OGCG: @_Z27assume_aligned_misalignmentPvj
+// OGCG:   call void @llvm.assume(i1 true) [ "align"(ptr %{{.+}}, i64 16, i64 %{{.+}}) ]
+// OGCG: }
+
 void assume_separate_storage(void *p1, void *p2) {
   __builtin_assume_separate_storage(p1, p2);
 }