[mlir][llvm] Fix bug in constant import from LLVM IR.

[gysit]

Summary

The revision addresses a bug during constant expression traversal
when importing LLVM IR. A constant expression may have cyclic
dependencies, for example, when a constant is initialized with its
address. This revision extends the constant expression traversal
to detect cyclic dependencies and adds a test to verify this
case is handled properly.

[gysit]

Switch to a hopefully easier to understand algorithm.

[Dinistro]

LGTM!
The algorithm change made it indeed more easy to understand.

mlir/lib/Target/LLVMIR/ModuleImport.cpp

ModuleImport::convertGlobalCtorsAndDtors(llvm::GlobalVariable *globalVar) {
   globalInsertionOp = builder.create<LLVM::GlobalDtorsOp>(
       mlirModule.getLoc(), builder.getArrayAttr(funcs),
       builder.getI32ArrayAttr(priorities));
   return success();
 }
 
 SetVector<llvm::Constant *>
 ModuleImport::getConstantsToConvert(llvm::Constant *constant) {
-  // Traverse the constant dependencies in post order.
-  SmallVector<llvm::Constant *> workList;
+  // Traverse the constants in pre-order and stop the traversal if a constant
+  // already has a `valueMapping` entry from an earlier constant translation or
+  // if a constant has a cyclic dependency.
   SmallVector<llvm::Constant *> orderedList;
-  workList.push_back(constant);
+  SmallVector<llvm::Constant *> subtreeSuccessorList;
+  DenseMap<llvm::Constant *, int64_t> constantToIndex;
+  SmallVector<llvm::Constant *> workList = {constant};
   while (!workList.empty()) {
     llvm::Constant *current = workList.pop_back_val();
-    // Skip constants that have been converted before and store all other ones.
+    // Skip constants that already have a `valueMapping` entry from an earlier
+    // constant translation.
     if (valueMapping.count(current))
       continue;
+    // Skip constants if a cyclic dependency is detect. A constant has a cyclic
+    // dependency if the subtree of its dependencies contains the constant
+    // itself. A constant that appears in distinct subtrees is not a cyclic
+    // dependency and has to be inserted multiple times into `orderedList`. The
+    // later insertion into `orderedSet` then keeps the duplicate that appears
+    // last during the traversal, which ensures all dependencies are satisfied.
+    if (constantToIndex.count(current)) {
+      // A cycle is detected if the subtree successor of the current constant,
+      // the constant that is traversed right after the current constant's
+      // subtree, has either not been traversed or has been traversed before the
+      // last traversal of the current constant. Note that the lookup of the
+      // subtree successor index returns zero if the successor has not been
+      // traversed or if there is no successor as in case of the root constant.
+      int64_t idx = constantToIndex.lookup(current);
+      if (constantToIndex.lookup(subtreeSuccessorList[idx]) <= idx)
+        continue;
+    }
+    // Add the current constant the the `orderedList` and store its index as
+    // well as the subtree successor at the top of `workList`. This successor
+    // will be processed next after traversing the current constant's subtree.
+    constantToIndex[current] = orderedList.size();
     orderedList.push_back(current);
+    subtreeSuccessorList.push_back(workList.empty() ? nullptr
+                                                    : workList.back());
     // Add the current constant's dependencies to the work list. Only add
     // constant dependencies and skip any other values such as basic block
     // addresses.
     for (llvm::Value *operand : current->operands())
       if (auto *constDependency = dyn_cast<llvm::Constant>(operand))
         workList.push_back(constDependency);
-    // Use the `getElementValue` method to add the dependencies of zero
+    // Use the getElementValue method to add the dependencies of zero
     // initialized aggregate constants since they do not take any operands.
     if (auto *constAgg = dyn_cast<llvm::ConstantAggregateZero>(current)) {
       unsigned numElements = constAgg->getElementCount().getFixedValue();
       for (unsigned i = 0, e = numElements; i != e; ++i)
         workList.push_back(constAgg->getElementValue(i));
     }
   }
 
-  // Add the constants in reverse post order to the result set to ensure all
+  // Add the constants in reverse pre-order to the result set to ensure all
   // dependencies are satisfied. Avoid storing duplicates since LLVM constants
   // are uniqued and only one `valueMapping` entry per constant is possible.
   SetVector<llvm::Constant *> orderedSet;
   for (llvm::Constant *orderedConst : llvm::reverse(orderedList))
     orderedSet.insert(orderedConst);
   return orderedSet;
 }
 

mlir/test/Target/LLVMIR/Import/constant.ll

define i64 @const_exprs_with_duplicate() {
   ; CHECK: llvm.getelementptr %[[ADDR]][%{{.*}}] : (!llvm.ptr, i32) -> !llvm.ptr
   %1 = add i64 1, ptrtoint (ptr getelementptr (i32, ptr @global, i32 7) to i64)
 
   ; Verify the address value is reused.
   ; CHECK: llvm.getelementptr %[[ADDR]][%{{.*}}] : (!llvm.ptr, i32) -> !llvm.ptr
   %2 = add i64 %1, ptrtoint (ptr getelementptr (i32, ptr @global, i32 42) to i64)
   ret i64 %2
 }
+
+; // -----
+
+; Verify the import of constant expressions with cyclic dependencies.
+
+@cyclic = internal constant i64 mul (i64 ptrtoint (ptr @cyclic to i64), i64 ptrtoint (ptr @cyclic to i64))
+
+; CHECK-LABEL: @cyclic
+; CHECK:  %[[ADDR:.+]] = llvm.mlir.addressof @cyclic
+; CHECK:  %[[VAL0:.+]] = llvm.ptrtoint %[[ADDR]]
+; CHECK:  %[[VAL1:.+]] = llvm.mul %[[VAL0]], %[[VAL0]]
+; CHECK:  llvm.return %[[VAL1]]