[mlir] Fix conflict of user defined reserved functions with internal prototypes #123378

Luohaothu · 2025-01-17T17:54:55Z

Related to #120950

On lowering from memref to LLVM, malloc and other intrinsic functions from libc will be declared in the current module. User's redefinition of these reserved functions will poison the internal analysis with wrong prototype. This patch adds assertion on the found function's type and reports if it mismatch with the intended type.

llvmbot · 2025-01-17T17:55:33Z

@llvm/pr-subscribers-mlir

@llvm/pr-subscribers-mlir-llvm

Author: Luohao Wang (Luohaothu)

Changes

Related to #120950

On lowering from memref to LLVM, malloc and other intrinsic functions from libc will be declared in the current module. User's redefinition of these reserved functions will poison the internal analysis with wrong prototype. This patch adds assertion on the found function's type and reports if it mismatch with the intended type.

Full diff: https://github.com/llvm/llvm-project/pull/123378.diff

3 Files Affected:

(modified) mlir/include/mlir/Dialect/LLVMIR/FunctionCallUtils.h (+2-1)
(modified) mlir/lib/Dialect/LLVMIR/IR/FunctionCallUtils.cpp (+36-20)
(added) mlir/test/Conversion/MemRefToLLVM/issue-120950.mlir (+11)

diff --git a/mlir/include/mlir/Dialect/LLVMIR/FunctionCallUtils.h b/mlir/include/mlir/Dialect/LLVMIR/FunctionCallUtils.h
index 852490cf7428f8..3095c83b90db9e 100644
--- a/mlir/include/mlir/Dialect/LLVMIR/FunctionCallUtils.h
+++ b/mlir/include/mlir/Dialect/LLVMIR/FunctionCallUtils.h
@@ -64,7 +64,8 @@ LLVM::LLVMFuncOp lookupOrCreateMemRefCopyFn(Operation *moduleOp, Type indexType,
 /// Create a FuncOp with signature `resultType`(`paramTypes`)` and name `name`.
 LLVM::LLVMFuncOp lookupOrCreateFn(Operation *moduleOp, StringRef name,
                                   ArrayRef<Type> paramTypes = {},
-                                  Type resultType = {}, bool isVarArg = false);
+                                  Type resultType = {}, bool isVarArg = false,
+                                  bool isReserved = false);
 
 } // namespace LLVM
 } // namespace mlir
diff --git a/mlir/lib/Dialect/LLVMIR/IR/FunctionCallUtils.cpp b/mlir/lib/Dialect/LLVMIR/IR/FunctionCallUtils.cpp
index 88421a16ccf9fb..ecc31df40ea52f 100644
--- a/mlir/lib/Dialect/LLVMIR/IR/FunctionCallUtils.cpp
+++ b/mlir/lib/Dialect/LLVMIR/IR/FunctionCallUtils.cpp
@@ -48,13 +48,29 @@ static constexpr llvm::StringRef kMemRefCopy = "memrefCopy";
 LLVM::LLVMFuncOp mlir::LLVM::lookupOrCreateFn(Operation *moduleOp,
                                               StringRef name,
                                               ArrayRef<Type> paramTypes,
-                                              Type resultType, bool isVarArg) {
+                                              Type resultType, bool isVarArg, bool isReserved) {
   assert(moduleOp->hasTrait<OpTrait::SymbolTable>() &&
          "expected SymbolTable operation");
   auto func = llvm::dyn_cast_or_null<LLVM::LLVMFuncOp>(
       SymbolTable::lookupSymbolIn(moduleOp, name));
-  if (func)
+  auto funcT = LLVMFunctionType::get(resultType, paramTypes, isVarArg);
+  // Assert the signature of the found function is same as expected
+  if (func) {
+    if (funcT != func.getFunctionType()) {
+      if (isReserved) {
+        func.emitError("redefinition of reserved function '" + name + "' of different type ")
+        .append(func.getFunctionType())
+        .append(" is prohibited");
+        exit(0);
+      } else {
+        func.emitError("redefinition of function '" + name + "' of different type ")
+        .append(funcT)
+        .append(" is prohibited");
+        exit(0);
+      }
+    }
     return func;
+  }
   OpBuilder b(moduleOp->getRegion(0));
   return b.create<LLVM::LLVMFuncOp>(
       moduleOp->getLoc(), name,
@@ -64,37 +80,37 @@ LLVM::LLVMFuncOp mlir::LLVM::lookupOrCreateFn(Operation *moduleOp,
 LLVM::LLVMFuncOp mlir::LLVM::lookupOrCreatePrintI64Fn(Operation *moduleOp) {
   return lookupOrCreateFn(moduleOp, kPrintI64,
                           IntegerType::get(moduleOp->getContext(), 64),
-                          LLVM::LLVMVoidType::get(moduleOp->getContext()));
+                          LLVM::LLVMVoidType::get(moduleOp->getContext()), false, true);
 }
 
 LLVM::LLVMFuncOp mlir::LLVM::lookupOrCreatePrintU64Fn(Operation *moduleOp) {
   return lookupOrCreateFn(moduleOp, kPrintU64,
                           IntegerType::get(moduleOp->getContext(), 64),
-                          LLVM::LLVMVoidType::get(moduleOp->getContext()));
+                          LLVM::LLVMVoidType::get(moduleOp->getContext()), false, true);
 }
 
 LLVM::LLVMFuncOp mlir::LLVM::lookupOrCreatePrintF16Fn(Operation *moduleOp) {
   return lookupOrCreateFn(moduleOp, kPrintF16,
                           IntegerType::get(moduleOp->getContext(), 16), // bits!
-                          LLVM::LLVMVoidType::get(moduleOp->getContext()));
+                          LLVM::LLVMVoidType::get(moduleOp->getContext()), false, true);
 }
 
 LLVM::LLVMFuncOp mlir::LLVM::lookupOrCreatePrintBF16Fn(Operation *moduleOp) {
   return lookupOrCreateFn(moduleOp, kPrintBF16,
                           IntegerType::get(moduleOp->getContext(), 16), // bits!
-                          LLVM::LLVMVoidType::get(moduleOp->getContext()));
+                          LLVM::LLVMVoidType::get(moduleOp->getContext()), false, true);
 }
 
 LLVM::LLVMFuncOp mlir::LLVM::lookupOrCreatePrintF32Fn(Operation *moduleOp) {
   return lookupOrCreateFn(moduleOp, kPrintF32,
                           Float32Type::get(moduleOp->getContext()),
-                          LLVM::LLVMVoidType::get(moduleOp->getContext()));
+                          LLVM::LLVMVoidType::get(moduleOp->getContext()), false, true);
 }
 
 LLVM::LLVMFuncOp mlir::LLVM::lookupOrCreatePrintF64Fn(Operation *moduleOp) {
   return lookupOrCreateFn(moduleOp, kPrintF64,
                           Float64Type::get(moduleOp->getContext()),
-                          LLVM::LLVMVoidType::get(moduleOp->getContext()));
+                          LLVM::LLVMVoidType::get(moduleOp->getContext()), false, true);
 }
 
 static LLVM::LLVMPointerType getCharPtr(MLIRContext *context) {
@@ -110,51 +126,51 @@ LLVM::LLVMFuncOp mlir::LLVM::lookupOrCreatePrintStringFn(
     Operation *moduleOp, std::optional<StringRef> runtimeFunctionName) {
   return lookupOrCreateFn(moduleOp, runtimeFunctionName.value_or(kPrintString),
                           getCharPtr(moduleOp->getContext()),
-                          LLVM::LLVMVoidType::get(moduleOp->getContext()));
+                          LLVM::LLVMVoidType::get(moduleOp->getContext()), false, true);
 }
 
 LLVM::LLVMFuncOp mlir::LLVM::lookupOrCreatePrintOpenFn(Operation *moduleOp) {
   return lookupOrCreateFn(moduleOp, kPrintOpen, {},
-                          LLVM::LLVMVoidType::get(moduleOp->getContext()));
+                          LLVM::LLVMVoidType::get(moduleOp->getContext()), false, true);
 }
 
 LLVM::LLVMFuncOp mlir::LLVM::lookupOrCreatePrintCloseFn(Operation *moduleOp) {
   return lookupOrCreateFn(moduleOp, kPrintClose, {},
-                          LLVM::LLVMVoidType::get(moduleOp->getContext()));
+                          LLVM::LLVMVoidType::get(moduleOp->getContext()), false, true);
 }
 
 LLVM::LLVMFuncOp mlir::LLVM::lookupOrCreatePrintCommaFn(Operation *moduleOp) {
   return lookupOrCreateFn(moduleOp, kPrintComma, {},
-                          LLVM::LLVMVoidType::get(moduleOp->getContext()));
+                          LLVM::LLVMVoidType::get(moduleOp->getContext()), false, true);
 }
 
 LLVM::LLVMFuncOp mlir::LLVM::lookupOrCreatePrintNewlineFn(Operation *moduleOp) {
   return lookupOrCreateFn(moduleOp, kPrintNewline, {},
-                          LLVM::LLVMVoidType::get(moduleOp->getContext()));
+                          LLVM::LLVMVoidType::get(moduleOp->getContext()), false, true);
 }
 
 LLVM::LLVMFuncOp mlir::LLVM::lookupOrCreateMallocFn(Operation *moduleOp,
                                                     Type indexType) {
   return LLVM::lookupOrCreateFn(moduleOp, kMalloc, indexType,
-                                getVoidPtr(moduleOp->getContext()));
+                                getVoidPtr(moduleOp->getContext()), false, true);
 }
 
 LLVM::LLVMFuncOp mlir::LLVM::lookupOrCreateAlignedAllocFn(Operation *moduleOp,
                                                           Type indexType) {
   return LLVM::lookupOrCreateFn(moduleOp, kAlignedAlloc, {indexType, indexType},
-                                getVoidPtr(moduleOp->getContext()));
+                                getVoidPtr(moduleOp->getContext()), false, true);
 }
 
 LLVM::LLVMFuncOp mlir::LLVM::lookupOrCreateFreeFn(Operation *moduleOp) {
   return LLVM::lookupOrCreateFn(
       moduleOp, kFree, getVoidPtr(moduleOp->getContext()),
-      LLVM::LLVMVoidType::get(moduleOp->getContext()));
+      LLVM::LLVMVoidType::get(moduleOp->getContext()), false, true);
 }
 
 LLVM::LLVMFuncOp mlir::LLVM::lookupOrCreateGenericAllocFn(Operation *moduleOp,
                                                           Type indexType) {
   return LLVM::lookupOrCreateFn(moduleOp, kGenericAlloc, indexType,
-                                getVoidPtr(moduleOp->getContext()));
+                                getVoidPtr(moduleOp->getContext()), false, true);
 }
 
 LLVM::LLVMFuncOp
@@ -162,13 +178,13 @@ mlir::LLVM::lookupOrCreateGenericAlignedAllocFn(Operation *moduleOp,
                                                 Type indexType) {
   return LLVM::lookupOrCreateFn(moduleOp, kGenericAlignedAlloc,
                                 {indexType, indexType},
-                                getVoidPtr(moduleOp->getContext()));
+                                getVoidPtr(moduleOp->getContext()), false, true);
 }
 
 LLVM::LLVMFuncOp mlir::LLVM::lookupOrCreateGenericFreeFn(Operation *moduleOp) {
   return LLVM::lookupOrCreateFn(
       moduleOp, kGenericFree, getVoidPtr(moduleOp->getContext()),
-      LLVM::LLVMVoidType::get(moduleOp->getContext()));
+      LLVM::LLVMVoidType::get(moduleOp->getContext()), false, true);
 }
 
 LLVM::LLVMFuncOp
@@ -177,5 +193,5 @@ mlir::LLVM::lookupOrCreateMemRefCopyFn(Operation *moduleOp, Type indexType,
   return LLVM::lookupOrCreateFn(
       moduleOp, kMemRefCopy,
       ArrayRef<Type>{indexType, unrankedDescriptorType, unrankedDescriptorType},
-      LLVM::LLVMVoidType::get(moduleOp->getContext()));
+      LLVM::LLVMVoidType::get(moduleOp->getContext()), false, true);
 }
diff --git a/mlir/test/Conversion/MemRefToLLVM/issue-120950.mlir b/mlir/test/Conversion/MemRefToLLVM/issue-120950.mlir
new file mode 100644
index 00000000000000..f744e4f7635ea7
--- /dev/null
+++ b/mlir/test/Conversion/MemRefToLLVM/issue-120950.mlir
@@ -0,0 +1,11 @@
+// RUN: mlir-opt %s -finalize-memref-to-llvm 2>&1 | FileCheck %s
+
+#map = affine_map<(d0) -> (d0 + 1)>
+module {
+  // CHECK: redefinition of reserved function 'malloc' of different type '!llvm.func<void (i64)>' is prohibited
+  llvm.func @malloc(i64)
+  func.func @issue_120950() {
+    %alloc = memref.alloc() : memref<1024x64xf32, 1>
+    llvm.return
+  }
+}

github-actions · 2025-01-17T18:03:18Z

✅ With the latest revision this PR passed the C/C++ code formatter.

Luohaothu · 2025-01-18T12:34:16Z

@wsmoses @jeanPerier can you help review this PR?

Luohaothu · 2025-01-20T13:21:21Z

Ping

Dinistro

Dropped a few comments. As described, using exit(0) is not how we deal with illegal inputs thus this should be changed.

mlir/lib/Dialect/LLVMIR/IR/FunctionCallUtils.cpp

mlir/test/Conversion/MemRefToLLVM/issue-120950.mlir

…or handling

Dinistro

Sorry for the delayed review.

Thank you very much for the improvements. I added one functional comment and a set of nits, but this is soon ready to be shipped 🙂

mlir/lib/Conversion/LLVMCommon/PrintCallHelper.cpp

Dinistro · 2025-01-24T09:48:17Z

mlir/lib/Conversion/MemRefToLLVM/AllocLikeConversion.cpp

      getTypeConverter(), op->getParentWithTrait<OpTrait::SymbolTable>(),
      getIndexType());
-  auto results = rewriter.create<LLVM::CallOp>(loc, allocFuncOp, sizeBytes);
+  if (failed(allocFuncOp))
+    return std::make_tuple(Value(), Value());


Will this result in a crash later on? If so, this would also require a FailureOr wrapping.

The returned pair of value is populated and finally verified here: AllocLikeConvertion.cpp

// Allocate the underlying buffer. auto [allocatedPtr, alignedPtr] = this->allocateBuffer(rewriter, loc, size, op); if (!allocatedPtr || !alignedPtr) return rewriter.notifyMatchFailure(loc, "underlying buffer allocation failed");

Here empty value is assigned as the invalid state and properly handled. I prefer to leave it as is, and a separate PR might be more appropriate to modernize the error handling style if we do have a preference.

Fine for me 🙂

Dinistro · 2025-01-24T09:48:23Z

mlir/lib/Conversion/MemRefToLLVM/AllocLikeConversion.cpp

      getTypeConverter(), op->getParentWithTrait<OpTrait::SymbolTable>(),
      getIndexType());
+  if (failed(allocFuncOp))
+    return Value();


Handled at MemRefToLLVM.cpp

Value ptr = allocateBufferAutoAlign( rewriter, loc, sizeBytes, op, &defaultLayout, alignedAllocationGetAlignment(rewriter, loc, cast<memref::AllocOp>(op), &defaultLayout)); if (!ptr) return std::make_tuple(Value(), Value()); return std::make_tuple(ptr, ptr);

in the same style.

mlir/lib/Conversion/MemRefToLLVM/MemRefToLLVM.cpp

mlir/lib/Conversion/VectorToLLVM/ConvertVectorToLLVM.cpp

mlir/lib/Dialect/LLVMIR/IR/FunctionCallUtils.cpp

mlir/lib/Conversion/MemRefToLLVM/MemRefToLLVM.cpp

mlir/lib/Conversion/MemRefToLLVM/AllocLikeConversion.cpp

Dinistro · 2025-01-24T09:56:48Z

mlir/test/Conversion/MemRefToLLVM/invalid.mlir

Nit: Can you change this test to use --split-input-file and --verify-diagnostics instead of using 2>&1? This makes the test more robust and modern.

The bad_address_space test emit error from unknown location, which cannot be handled by --verify-diagnostics. Do you have any suggestions?

Feel free to keep this as is. If this should change, the owners of that test should do so.

Dinistro

LGTM, thanks for addressing the comments.

Maybe give people another day to also take a look, then you can land this.

Dinistro · 2025-01-27T09:34:23Z

mlir/lib/Conversion/MemRefToLLVM/AllocLikeConversion.cpp

      getTypeConverter(), op->getParentWithTrait<OpTrait::SymbolTable>(),
      getIndexType());
-  auto results = rewriter.create<LLVM::CallOp>(loc, allocFuncOp, sizeBytes);
+  if (failed(allocFuncOp))
+    return std::make_tuple(Value(), Value());


Fine for me 🙂

Dinistro · 2025-01-27T09:36:47Z

mlir/test/Conversion/MemRefToLLVM/invalid.mlir

Feel free to keep this as is. If this should change, the owners of that test should do so.

Luohaothu · 2025-01-27T10:35:49Z

Big thanks for the review @Dinistro ! I can hold it for one day or two. Could you help merge the PR after that?

By the way, how to gain write access for the repo? Currently I have to ask for help by comment each time😓

jeanPerier

This is outside of my area, but the patch makes sense to me. We did hit similar crashes in flang (like here), and having proper asserts for helpers retrieving libc functions seems right to me.

jeanPerier · 2025-01-27T13:01:07Z

By the way, how to gain write access for the repo? Currently I have to ask for help by comment each time😓

You have to open an issue with special tags to request commit access, see https://llvm.org/docs/DeveloperPolicy.html#obtaining-commit-access

Dinistro · 2025-01-27T13:43:46Z

Big thanks for the review @Dinistro ! I can hold it for one day or two. Could you help merge the PR after that?

By the way, how to gain write access for the repo? Currently I have to ask for help by comment each time😓

I'll add a reminder to land this tomorrow. In case I forget, feel free to ping me 😅

llvmbot added mlir:llvm mlir labels Jan 17, 2025

Dinistro self-requested a review January 20, 2025 15:00

Dinistro reviewed Jan 20, 2025

View reviewed changes

mlir/lib/Dialect/LLVMIR/IR/FunctionCallUtils.cpp Outdated Show resolved Hide resolved

mlir/lib/Dialect/LLVMIR/IR/FunctionCallUtils.cpp Outdated Show resolved Hide resolved

mlir/test/Conversion/MemRefToLLVM/issue-120950.mlir Outdated Show resolved Hide resolved

Luohaothu force-pushed the res-fix branch from 832a33e to 11720e3 Compare January 21, 2025 07:39

Luohaothu requested review from banach-space, dcaballe and nicolasvasilache as code owners January 21, 2025 10:17

Luohaothu and others added 5 commits January 21, 2025 18:17

[mlir] Add assertion on reserved function's type

3b892e6

[mlir] Add test

bdb69fc

[mlir] Reformat code

d26a77d

[mlir] Wrapped return value of function lookup in FailureOr for err…

5ec38d6

…or handling

[mlir] [Test] Moved & renamed test case

9079caf

Luohaothu force-pushed the res-fix branch from 93e19aa to 9079caf Compare January 21, 2025 10:18

Reformat code

475409c

Luohaothu requested review from Dinistro and jeanPerier January 21, 2025 10:34

Dinistro reviewed Jan 24, 2025

View reviewed changes

Make stylish fixes

874a1cd

Luohaothu requested a review from Dinistro January 26, 2025 07:16

Luohaothu added 2 commits January 26, 2025 15:22

More stylish fixes

2d7dc5d

Format code

a7f3308

Dinistro approved these changes Jan 27, 2025

View reviewed changes

jeanPerier reviewed Jan 27, 2025

View reviewed changes

Luohaothu mentioned this pull request Jan 27, 2025

Request Commit Access For luohaothu #124564

Closed

Dinistro merged commit e84f6b6 into llvm:main Jan 28, 2025
8 checks passed

Luohaothu deleted the res-fix branch February 1, 2025 07:24

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[mlir] Fix conflict of user defined reserved functions with internal prototypes #123378

[mlir] Fix conflict of user defined reserved functions with internal prototypes #123378

Luohaothu commented Jan 17, 2025

llvmbot commented Jan 17, 2025 •

edited

Loading

github-actions bot commented Jan 17, 2025 •

edited

Loading

Luohaothu commented Jan 18, 2025

Luohaothu commented Jan 20, 2025

Dinistro left a comment

Dinistro left a comment

Dinistro Jan 24, 2025

Luohaothu Jan 26, 2025

Dinistro Jan 27, 2025

Dinistro Jan 24, 2025

Luohaothu Jan 26, 2025

Dinistro Jan 24, 2025

Luohaothu Jan 26, 2025

Dinistro Jan 27, 2025

Dinistro left a comment

Dinistro Jan 27, 2025

Dinistro Jan 27, 2025

Luohaothu commented Jan 27, 2025

jeanPerier left a comment

jeanPerier commented Jan 27, 2025

Dinistro commented Jan 27, 2025

[mlir] Fix conflict of user defined reserved functions with internal prototypes #123378

[mlir] Fix conflict of user defined reserved functions with internal prototypes #123378

Conversation

Luohaothu commented Jan 17, 2025

llvmbot commented Jan 17, 2025 • edited Loading

github-actions bot commented Jan 17, 2025 • edited Loading

Luohaothu commented Jan 18, 2025

Luohaothu commented Jan 20, 2025

Dinistro left a comment

Choose a reason for hiding this comment

Dinistro left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Dinistro left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Luohaothu commented Jan 27, 2025

jeanPerier left a comment

Choose a reason for hiding this comment

jeanPerier commented Jan 27, 2025

Dinistro commented Jan 27, 2025

llvmbot commented Jan 17, 2025 •

edited

Loading

github-actions bot commented Jan 17, 2025 •

edited

Loading