[CIR][CIRGen] Support __builtin___memset_chk (#1053)

This follows the same implementation as CodeGen. #1051 tracks potentially switching to a different strategy in the future.
llvm · Nov 4, 2024 · 815dfaa · 815dfaa
1 parent d4f09e6
commit 815dfaa
Show file tree

Hide file tree

Showing 2 changed files with 70 additions and 2 deletions.
diff --git a/clang/lib/CIR/CodeGen/CIRGenBuiltin.cpp b/clang/lib/CIR/CodeGen/CIRGenBuiltin.cpp
@@ -1478,8 +1478,23 @@ RValue CIRGenFunction::buildBuiltinExpr(const GlobalDecl GD, unsigned BuiltinID,
 
   case Builtin::BI__builtin_memset_inline:
     llvm_unreachable("BI__builtin_memset_inline NYI");
-  case Builtin::BI__builtin___memset_chk:
-    llvm_unreachable("BI__builtin___memset_chk NYI");
+  case Builtin::BI__builtin___memset_chk: {
+    // fold __builtin_memset_chk(x, y, cst1, cst2) to memset iff cst1<=cst2.
+    Expr::EvalResult sizeResult, dstSizeResult;
+    if (!E->getArg(2)->EvaluateAsInt(sizeResult, CGM.getASTContext()) ||
+        !E->getArg(3)->EvaluateAsInt(dstSizeResult, CGM.getASTContext()))
+      break;
+    llvm::APSInt size = sizeResult.Val.getInt();
+    llvm::APSInt dstSize = dstSizeResult.Val.getInt();
+    if (size.ugt(dstSize))
+      break;
+    Address dest = buildPointerWithAlignment(E->getArg(0));
+    mlir::Value byteVal = buildScalarExpr(E->getArg(1));
+    auto loc = getLoc(E->getSourceRange());
+    ConstantOp sizeOp = builder.getConstInt(loc, size);
+    builder.createMemSet(loc, dest.getPointer(), byteVal, sizeOp);
+    return RValue::get(dest.getPointer());
+  }
   case Builtin::BI__builtin_wmemchr:
     llvm_unreachable("BI__builtin_wmemchr NYI");
   case Builtin::BI__builtin_wmemcmp:

diff --git a/clang/test/CIR/CodeGen/builtins-memory.c b/clang/test/CIR/CodeGen/builtins-memory.c
@@ -54,3 +54,56 @@ void test_memcpy_chk(void *dest, const void *src, size_t n) {
   // CIR: cir.call @__memcpy_chk(%[[#DEST_LOAD]], %[[#SRC_LOAD]], %[[#N_LOAD1]], %[[#N_LOAD2]])
   __builtin___memcpy_chk(dest, src, n, n);
 }
+
+void test_memset_chk(void *dest, int ch, size_t n) {
+  // CIR-LABEL: cir.func @test_memset_chk
+  // CIR:         %[[#DEST:]] = cir.alloca {{.*}} ["dest", init]
+  // CIR:         %[[#CH:]] = cir.alloca {{.*}} ["ch", init]
+  // CIR:         %[[#N:]] = cir.alloca {{.*}} ["n", init]
+
+  // An unchecked memset should be emitted when the count and buffer size are
+  // constants and the count is less than or equal to the buffer size.
+
+  // CIR: %[[#DEST_LOAD:]] = cir.load %[[#DEST]]
+  // CIR: %[[#CH_LOAD:]] = cir.load %[[#CH]]
+  // CIR: %[[#COUNT:]] = cir.const #cir.int<8>
+  // CIR: cir.libc.memset %[[#COUNT]] bytes from %[[#DEST_LOAD]] set to %[[#CH_LOAD]]
+  __builtin___memset_chk(dest, ch, 8, 10);
+
+  // CIR: %[[#DEST_LOAD:]] = cir.load %[[#DEST]]
+  // CIR: %[[#CH_LOAD:]] = cir.load %[[#CH]]
+  // CIR: %[[#COUNT:]] = cir.const #cir.int<10>
+  // CIR: cir.libc.memset %[[#COUNT]] bytes from %[[#DEST_LOAD]] set to %[[#CH_LOAD]]
+  __builtin___memset_chk(dest, ch, 10, 10);
+
+  // __memset_chk should be called when the count is greater than the buffer
+  // size, or when either the count or buffer size isn't a constant.
+
+  // CIR: %[[#DEST_LOAD:]] = cir.load %[[#DEST]]
+  // CIR: %[[#CH_LOAD:]] = cir.load %[[#CH]]
+  // CIR: %[[#COUNT:]] = cir.const #cir.int<10>
+  // CIR: %[[#SIZE:]] = cir.const #cir.int<8>
+  // CIR: cir.call @__memset_chk(%[[#DEST_LOAD]], %[[#CH_LOAD]], %[[#COUNT]], %[[#SIZE]])
+  __builtin___memset_chk(dest, ch, 10lu, 8lu);
+
+  // CIR: %[[#DEST_LOAD:]] = cir.load %[[#DEST]]
+  // CIR: %[[#CH_LOAD:]] = cir.load %[[#CH]]
+  // CIR: %[[#N_LOAD:]] = cir.load %[[#N]]
+  // CIR: %[[#SIZE:]] = cir.const #cir.int<10>
+  // CIR: cir.call @__memset_chk(%[[#DEST_LOAD]], %[[#CH_LOAD]], %[[#N_LOAD]], %[[#SIZE]])
+  __builtin___memset_chk(dest, ch, n, 10lu);
+
+  // CIR: %[[#DEST_LOAD:]] = cir.load %[[#DEST]]
+  // CIR: %[[#CH_LOAD:]] = cir.load %[[#CH]]
+  // CIR: %[[#COUNT:]] = cir.const #cir.int<10>
+  // CIR: %[[#N_LOAD:]] = cir.load %[[#N]]
+  // CIR: cir.call @__memset_chk(%[[#DEST_LOAD]], %[[#CH_LOAD]], %[[#COUNT]], %[[#N_LOAD]])
+  __builtin___memset_chk(dest, ch, 10lu, n);
+
+  // CIR: %[[#DEST_LOAD:]] = cir.load %[[#DEST]]
+  // CIR: %[[#CH_LOAD:]] = cir.load %[[#CH]]
+  // CIR: %[[#N_LOAD1:]] = cir.load %[[#N]]
+  // CIR: %[[#N_LOAD2:]] = cir.load %[[#N]]
+  // CIR: cir.call @__memset_chk(%[[#DEST_LOAD]], %[[#CH_LOAD]], %[[#N_LOAD1]], %[[#N_LOAD2]])
+  __builtin___memset_chk(dest, ch, n, n);
+}