[mlir][LLVMIR] Do not update instMap via assignments to map entry references

[myhsu]

Summary

Inside processInstruction, we assign the translated mlir::Value to a reference previously taken from the corresponding entry in instMap. However, instMap (a DenseMap) might resize after the entry reference was taken, rendering the assignment useless since it's assigning to a dangling reference. Here is a (pseudo) snippet that shows the concept:

// inst has type llvm::Instruction *
Value &v = instMap[inst];
...
// op is one of the operands of inst, has type llvm::Value *
processValue(op);
// instMap resizes inside processValue
...
translatedValue = b.createOp<Foo>(...);
// v is already a dangling reference at this point!
// The following assignment is bogus.
v = translatedValue;

Nevertheless, after we stop caching llvm::Constant into instMap, there is only one case that can cause processValue to resize instMap: If the operand is a llvm::ConstantExpr. In which case we will insert the derived llvm::Instruction into instMap. To trigger instMap to resize, which is a DenseMap, the threshold depends on the ratio between # of map entries and # of (hash) buckets. More specifically, it resizes if (# of map entries / # of buckets) >= 0.75. In this case # of map entries is equal to # of LLVM instructions, and # of buckets is the power-of two upperbound of # of map entries. Thus, eventually in the attaching test case (test/Target/LLVMIR/Import/incorrect-instmap-assignment.ll), we picked 96 and 128 for the # of map entries and # of buckets, respectively. (We can't pick numbers that are too small since DenseMap used inlined storage for small number of entries). Therefore, the ConstantExpr in the
said test case (i.e. a GEP) is the 96-th llvm::Value cached into the instMap, triggering the issue we're discussing here on its enclosing instruction (i.e. a load).

This patch fixes this issue by calling operator[] everytime we need to
update an entry.

[Mogball]

mlir/lib/Target/LLVMIR/ConvertFromLLVMIR.cpp

737–738

Can the variable itself just be deleted?

[myhsu]

Remove unused variable v and update the assertion check.

[myhsu]

mlir/lib/Target/LLVMIR/ConvertFromLLVMIR.cpp

737–738

Yes, I'd removed this line and replaced the check in next line with !instMap.count(inst)

[Mogball]

lgtm but you might want to wait if anyone else has complaints

[ftynse]

The complexity of the test is scary. Is it just a matter of having a sufficient amount of instructions and/or constants in a vector?

[myhsu]

In D124627#3482001, @ftynse wrote:

The complexity of the test is scary. Is it just a matter of having a sufficient amount of instructions and/or constants in a vector?

Yeah I agree. I think it's just having sufficient number of instructions (since we're not caching constants into instMap anymore) to saturate instMap and force it to resize. I'll try to simplify the test, primarily removing some scary global variables/types.

[myhsu]

Update the test case.

[myhsu]

In D124627#3482001, @ftynse wrote:

The complexity of the test is scary. Is it just a matter of having a sufficient amount of instructions and/or constants in a vector?

It turns out it's much more tricky to trigger the issue. Please checkout the updated summary above for the rationale.

mlir/lib/Target/LLVMIR/ConvertFromLLVMIR.cpp

for (auto inst = target->begin(); isa<llvm::PHINode>(inst); ++inst) {
     blockArguments.push_back(value);
   }
   return success();
 }
 
 LogicalResult Importer::processInstruction(llvm::Instruction *inst) {
   // FIXME: Support uses of SubtargetData. Currently inbounds GEPs, fast-math
   // flags and call / operand attributes are not supported.
   Location loc = processDebugLoc(inst->getDebugLoc(), inst);
-  Value &v = instMap[inst];
-  assert(!v && "processInstruction must be called only once per instruction!");
+  assert(!instMap.count(inst) &&
+         "processInstruction must be called only once per instruction!");
   switch (inst->getOpcode()) {
   default:
     return emitError(loc) << "unknown instruction: " << diag(*inst);
   case llvm::Instruction::Add:
   case llvm::Instruction::FAdd:
   case llvm::Instruction::Sub:
   case llvm::Instruction::FSub:
   case llvm::Instruction::Mul:
case llvm::Instruction::BitCast: {
     if (!inst->getType()->isVoidTy()) {
       Type type = processType(inst->getType());
       if (!type)
         return failure();
       state.addTypes(type);
     }
     Operation *op = b.create(state);
     if (!inst->getType()->isVoidTy())
-      v = op->getResult(0);
+      instMap[inst] = op->getResult(0);
     return success();
   }
   case llvm::Instruction::Alloca: {
     Value size = processValue(inst->getOperand(0));
     if (!size)
       return failure();
 
     auto *allocaInst = cast<llvm::AllocaInst>(inst);
-    v = b.create<AllocaOp>(loc, processType(inst->getType()),
+    instMap[inst] =
+        b.create<AllocaOp>(loc, processType(inst->getType()),
                            processType(allocaInst->getAllocatedType()), size,
                            allocaInst->getAlign().value());
     return success();
   }
   case llvm::Instruction::ICmp: {
     Value lhs = processValue(inst->getOperand(0));
     Value rhs = processValue(inst->getOperand(1));
     if (!lhs || !rhs)
       return failure();
-    v = b.create<ICmpOp>(
+    instMap[inst] = b.create<ICmpOp>(
         loc, getICmpPredicate(cast<llvm::ICmpInst>(inst)->getPredicate()), lhs,
         rhs);
     return success();
   }
   case llvm::Instruction::FCmp: {
     Value lhs = processValue(inst->getOperand(0));
     Value rhs = processValue(inst->getOperand(1));
     if (!lhs || !rhs)
case llvm::Instruction::Switch: {
     b.create<SwitchOp>(loc, condition, blocks[defaultBB], defaultBlockArgs,
                        caseValues, caseBlocks, caseOperandRefs);
     return success();
   }
   case llvm::Instruction::PHI: {
     Type type = processType(inst->getType());
     if (!type)
       return failure();
-    v = b.getInsertionBlock()->addArgument(
+    instMap[inst] = b.getInsertionBlock()->addArgument(
         type, processDebugLoc(inst->getDebugLoc(), inst));
     return success();
   }
   case llvm::Instruction::Call: {
     llvm::CallInst *ci = cast<llvm::CallInst>(inst);
     SmallVector<Value, 4> ops;
     ops.reserve(inst->getNumOperands());
     for (auto &op : ci->args()) {
case llvm::Instruction::Call: {
     } else {
       Value calledValue = processValue(ci->getCalledOperand());
       if (!calledValue)
         return failure();
       ops.insert(ops.begin(), calledValue);
       op = b.create<CallOp>(loc, tys, ops);
     }
     if (!ci->getType()->isVoidTy())
-      v = op->getResult(0);
+      instMap[inst] = op->getResult(0);
     return success();
   }
   case llvm::Instruction::LandingPad: {
     llvm::LandingPadInst *lpi = cast<llvm::LandingPadInst>(inst);
     SmallVector<Value, 4> ops;
 
     for (unsigned i = 0, ie = lpi->getNumClauses(); i < ie; i++)
       ops.push_back(processConstant(lpi->getClause(i)));
 
     Type ty = processType(lpi->getType());
     if (!ty)
       return failure();
 
-    v = b.create<LandingpadOp>(loc, ty, lpi->isCleanup(), ops);
+    instMap[inst] = b.create<LandingpadOp>(loc, ty, lpi->isCleanup(), ops);
     return success();
   }
   case llvm::Instruction::Invoke: {
     llvm::InvokeInst *ii = cast<llvm::InvokeInst>(inst);
 
     SmallVector<Type, 2> tys;
     if (!ii->getType()->isVoidTy())
       tys.push_back(processType(inst->getType()));
case llvm::Instruction::Invoke: {
     } else {
       ops.insert(ops.begin(), processValue(ii->getCalledOperand()));
       op = b.create<InvokeOp>(loc, tys, ops, blocks[ii->getNormalDest()],
                               normalArgs, blocks[ii->getUnwindDest()],
                               unwindArgs);
     }
 
     if (!ii->getType()->isVoidTy())
-      v = op->getResult(0);
+      instMap[inst] = op->getResult(0);
     return success();
   }
   case llvm::Instruction::Fence: {
     StringRef syncscope;
     SmallVector<StringRef, 4> ssNs;
     llvm::LLVMContext &llvmContext = inst->getContext();
     llvm::FenceInst *fence = cast<llvm::FenceInst>(inst);
     llvmContext.getSyncScopeNames(ssNs);
for (const auto &en :
         staticIndices.push_back(
             static_cast<int32_t>(constantInt->getValue().getZExtValue()));
       }
     }
 
     Type type = processType(inst->getType());
     if (!type)
       return failure();
-    v = b.create<GEPOp>(loc, type, sourceElementType, basePtr, dynamicIndices,
-                        staticIndices);
+    instMap[inst] = b.create<GEPOp>(loc, type, sourceElementType, basePtr,
+                                    dynamicIndices, staticIndices);
     return success();
   }
   }
 }
 
 FlatSymbolRefAttr Importer::getPersonalityAsAttr(llvm::Function *f) {
   if (!f->hasPersonalityFn())
     return nullptr;

mlir/test/Target/LLVMIR/Import/incorrect-instmap-assignment.ll

+; RUN: mlir-translate --import-llvm %s | FileCheck %s
+
+; This test file is meant to saturate `instMap` used in the translation
+; and force it to resize.
+
+; This test is primarily used to make sure it doesn't bail out with non-zero
+; exit code. Thus, we only wrote minimum level of checks.
+
+%my_struct = type {i32, i32}
+@gvar = external global %my_struct
+
+; CHECK: llvm.func @f(%arg0: i32, %arg1: i32)
+define void @f(i32 %0, i32 %1) {
+  %3 = add i32 %0, %1
+  %4 = add i32 %1, %3
+  %5 = add i32 %3, %4
+  %6 = add i32 %4, %5
+  %7 = add i32 %5, %6
+  %8 = add i32 %6, %7
+  %9 = add i32 %7, %8
+  %10 = add i32 %8, %9
+  %11 = add i32 %9, %10
+  %12 = add i32 %10, %11
+  %13 = add i32 %11, %12
+  %14 = add i32 %12, %13
+  %15 = add i32 %13, %14
+  %16 = add i32 %14, %15
+  %17 = add i32 %15, %16
+  %18 = add i32 %16, %17
+  %19 = add i32 %17, %18
+  %20 = add i32 %18, %19
+  %21 = add i32 %19, %20
+  %22 = add i32 %20, %21
+  %23 = add i32 %21, %22
+  %24 = add i32 %22, %23
+  %25 = add i32 %23, %24
+  %26 = add i32 %24, %25
+  %27 = add i32 %25, %26
+  %28 = add i32 %26, %27
+  %29 = add i32 %27, %28
+  %30 = add i32 %28, %29
+  %31 = add i32 %29, %30
+  %32 = add i32 %30, %31
+  %33 = add i32 %31, %32
+  %34 = add i32 %32, %33
+  %35 = add i32 %33, %34
+  %36 = add i32 %34, %35
+  %37 = add i32 %35, %36
+  %38 = add i32 %36, %37
+  %39 = add i32 %37, %38
+  %40 = add i32 %38, %39
+  %41 = add i32 %39, %40
+  %42 = add i32 %40, %41
+  %43 = add i32 %41, %42
+  %44 = add i32 %42, %43
+  %45 = add i32 %43, %44
+  %46 = add i32 %44, %45
+  %47 = add i32 %45, %46
+  %48 = add i32 %46, %47
+  %49 = add i32 %47, %48
+  %50 = add i32 %48, %49
+  %51 = add i32 %49, %50
+  %52 = add i32 %50, %51
+  %53 = add i32 %51, %52
+  %54 = add i32 %52, %53
+  %55 = add i32 %53, %54
+  %56 = add i32 %54, %55
+  %57 = add i32 %55, %56
+  %58 = add i32 %56, %57
+  %59 = add i32 %57, %58
+  %60 = add i32 %58, %59
+  %61 = add i32 %59, %60
+  %62 = add i32 %60, %61
+  %63 = add i32 %61, %62
+  %64 = add i32 %62, %63
+  %65 = add i32 %63, %64
+  %66 = add i32 %64, %65
+  %67 = add i32 %65, %66
+  %68 = add i32 %66, %67
+  %69 = add i32 %67, %68
+  %70 = add i32 %68, %69
+  %71 = add i32 %69, %70
+  %72 = add i32 %70, %71
+  %73 = add i32 %71, %72
+  %74 = add i32 %72, %73
+  %75 = add i32 %73, %74
+  %76 = add i32 %74, %75
+  %77 = add i32 %75, %76
+  %78 = add i32 %76, %77
+  %79 = add i32 %77, %78
+  %80 = add i32 %78, %79
+  %81 = add i32 %79, %80
+  %82 = add i32 %80, %81
+  %83 = add i32 %81, %82
+  %84 = add i32 %82, %83
+  %85 = add i32 %83, %84
+  %86 = add i32 %84, %85
+  %87 = add i32 %85, %86
+  %88 = add i32 %86, %87
+  %89 = add i32 %87, %88
+  %90 = add i32 %88, %89
+  %91 = add i32 %89, %90
+  %92 = add i32 %90, %91
+  %93 = add i32 %91, %92
+  %94 = add i32 %92, %93
+  %95 = load i32, i32* getelementptr inbounds (%my_struct, %my_struct* @gvar, i32 0, i32 0)
+  %96 = add i32 %1, %95
+  ret void
+}