[mlir][SCCP] Track all source dialects when meeting two lattice values

[rriddle]

Summary

Some dialects have implicit casts that change the type of an input. For example, some call operations perform an implicit cast(such as a shape refinement) when passing arguments and results. This revision allows for tracking all of the source dialects so that constants can still be materialized properly in these situations.

[mehdi_amini]

mlir/lib/Transforms/SCCP.cpp

30

31

300

This comment is confusing to me:

• ownerDialect does not match a parameter
• **the** source dialect of 'to' or 'from' : this revision is updating LatticeValue to have a set of dialects, also it isn't clear what this means in the context of the caller (it does not help me to know how to call this API properly)
• "'to' corresponds to a value whose uses should be visited if the meet changed" => I think you meant "value" instead of "to" here?

mlir/lib/Dialect/StandardOps/IR/Ops.cpp

#include "mlir/Dialect/StandardOps/IR/Ops.cpp.inc"
   addInterfaces<StdInlinerInterface>();
 }
 
 /// Materialize a single constant operation from a given attribute value with
 /// the desired resultant type.
 Operation *StandardOpsDialect::materializeConstant(OpBuilder &builder,
                                                    Attribute value, Type type,
                                                    Location loc) {
+  if (ConstantOp::isBuildableWith(value, type))
     return builder.create<ConstantOp>(loc, type, value);
+  return nullptr;
 }
 
 void mlir::printDimAndSymbolList(Operation::operand_iterator begin,
                                  Operation::operand_iterator end,
                                  unsigned numDims, OpAsmPrinter &p) {
   Operation::operand_range operands(begin, end);
   p << '(' << operands.take_front(numDims) << ')';
   if (operands.size() != numDims)

mlir/lib/Transforms/SCCP.cpp

 #include "PassDetail.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/Dialect.h"
 #include "mlir/Interfaces/ControlFlowInterfaces.h"
 #include "mlir/Interfaces/SideEffects.h"
 #include "mlir/Pass/Pass.h"
 #include "mlir/Transforms/FoldUtils.h"
 #include "mlir/Transforms/Passes.h"
+#include "llvm/ADT/SetVector.h"
 
 using namespace mlir;
 
 namespace {
 /// This class represents a single lattice value. A lattive value corresponds to
 /// the various different states that a value in the SCCP dataflow anaylsis can
 /// take. See 'Kind' below for more details on the different states a value can
 /// take.
 class LatticeValue {
   enum Kind {
     /// A value with a yet to be determined value. This state may be changed to
     /// anything.
     Unknown,
 
     /// A value that is known to be a constant. This state may be changed to
     /// overdefined.
     Constant,
 
     /// A value that cannot statically be determined to be a constant. This
     /// state cannot be changed.
     Overdefined
   };
 
 public:
   /// Initialize a lattice value with "Unknown".
-  LatticeValue()
-      : constantAndTag(nullptr, Kind::Unknown), constantDialect(nullptr) {}
+  LatticeValue() : constantAndTag(nullptr, Kind::Unknown) {}
   /// Initialize a lattice value with a constant.
   LatticeValue(Attribute attr, Dialect *dialect)
-      : constantAndTag(attr, Kind::Constant), constantDialect(dialect) {}
+      : constantAndTag(attr, Kind::Constant) {
+    constantDialects.insert(dialect);
+  }
 
   /// Returns true if this lattice value is unknown.
   bool isUnknown() const { return constantAndTag.getInt() == Kind::Unknown; }
 
   /// Mark the lattice value as overdefined.
   void markOverdefined() {
     constantAndTag.setPointerAndInt(nullptr, Kind::Overdefined);
-    constantDialect = nullptr;
+    constantDialects.clear();
   }
 
   /// Returns true if the lattice is overdefined.
   bool isOverdefined() const {
     return constantAndTag.getInt() == Kind::Overdefined;
   }
 
   /// Mark the lattice value as constant.
   void markConstant(Attribute value, Dialect *dialect) {
     constantAndTag.setPointerAndInt(value, Kind::Constant);
-    constantDialect = dialect;
+    constantDialects.insert(dialect);
   }
 
   /// If this lattice is constant, return the constant. Returns nullptr
   /// otherwise.
   Attribute getConstant() const { return constantAndTag.getPointer(); }
 
   /// If this lattice is constant, return the dialect to use when materializing
   /// the constant.
-  Dialect *getConstantDialect() const {
+  ArrayRef<Dialect *> getConstantDialects() const {
     assert(getConstant() && "expected valid constant");
-    return constantDialect;
+    return constantDialects.getArrayRef();
+  }
+
+  /// Add the given dialect to the set of constant source dialects.
+  void addConstantDialect(Dialect *dialect) {
+    assert(getConstant() && "expected valid constant");
+    constantDialects.insert(dialect);
   }
 
   /// Merge in the value of the 'rhs' lattice into this one. Returns true if the
   /// lattice value changed.
   bool meet(const LatticeValue &rhs) {
     // If we are already overdefined, or rhs is unknown, there is nothing to do.
     if (isOverdefined() || rhs.isUnknown())
       return false;
     // If we are unknown, just take the value of rhs.
     if (isUnknown()) {
       constantAndTag = rhs.constantAndTag;
-      constantDialect = rhs.constantDialect;
+      constantDialects.set_union(rhs.constantDialects);
       return true;
     }
 
     // Otherwise, if this value doesn't match rhs go straight to overdefined.
     if (constantAndTag != rhs.constantAndTag) {
       markOverdefined();
       return true;
     }
+    constantDialects.set_union(rhs.constantDialects);
     return false;
   }
 
 private:
   /// The attribute value if this is a constant and the tag for the element
   /// kind.
   llvm::PointerIntPair<Attribute, 2, Kind> constantAndTag;
 
-  /// The dialect the constant originated from. This is only valid if the
+  /// The dialects the constant originated from. This is only valid if the
   /// lattice is a constant. This is not used as part of the key, and is only
   /// needed to materialize the held constant if necessary.
-  Dialect *constantDialect;
+  llvm::SetVector<Dialect *, SmallVector<Dialect *, 1>,
+                  SmallPtrSet<Dialect *, 1>>
+      constantDialects;
 };
 
 /// This class contains various state used when computing the lattice of a
 /// callable operation.
 class CallableLatticeState {
 public:
   /// Build a lattice state with a given callable region, and a specified number
   /// of results to be initialized to the default lattice value (Unknown).
for (auto value : values) {
         visitUsers(value);
       }
     }
   }
 
   /// Returns true if the given value was marked as overdefined.
   bool isOverdefined(Value value) const;
 
-  /// Merge in the given lattice 'from' into the lattice 'to'. 'owner'
-  /// corresponds to the parent operation of 'to'.
-  void meet(Operation *owner, LatticeValue &to, const LatticeValue &from);
+  /// Merge in the given lattice 'from' into the lattice 'to'. 'op' corresponds
+  /// to the source operation of either 'to' or 'from' that should be added back
+  /// to the worklist if the meet results in a change.
+  void meet(Operation *op, LatticeValue &to, const LatticeValue &from);
+
+  /// Merge in the given lattice 'from' into the lattice 'to'. 'to' corresponds
+  /// to a value whose uses should be visited if the meet changed.
+  /// 'ownerDialect' corresponds to the source dialect of 'to' or 'from'. This
+  /// is added to 'to' if it resolves to a constant.
+  void meet(Value value, Dialect *dialect, LatticeValue &to,
+            const LatticeValue &from);
 
   /// The lattice for each SSA value.
   DenseMap<Value, LatticeValue> latticeValues;
 
   /// The set of blocks that are known to execute, or are intrinsically live.
   SmallPtrSet<Block *, 16> executableBlocks;
 
   /// The set of control flow edges that are known to execute.
LogicalResult SCCPSolver::replaceWithConstant(OpBuilder &builder,
                                               OperationFolder &folder,
                                               Value value) {
   auto it = latticeValues.find(value);
   auto attr = it == latticeValues.end() ? nullptr : it->second.getConstant();
   if (!attr)
     return failure();
 
   // Attempt to materialize a constant for the given value.
-  Dialect *dialect = it->second.getConstantDialect();
-  Value constant = folder.getOrCreateConstant(builder, dialect, attr,
-                                              value.getType(), value.getLoc());
-  if (!constant)
-    return failure();
-
+  Value constant;
+  for (Dialect *dialect : it->second.getConstantDialects()) {
+    if (Value constant = folder.getOrCreateConstant(
+            builder, dialect, attr, value.getType(), value.getLoc())) {
       value.replaceAllUsesWith(constant);
       latticeValues.erase(it);
       return success();
     }
+  }
+  return failure();
+}
 
 void SCCPSolver::visitOperation(Operation *op) {
   // Collect all of the constant operands feeding into this operation. If any
   // are not ready to be resolved, bail out and wait for them to resolve.
   SmallVector<Attribute, 8> operandConstants;
   operandConstants.reserve(op->getNumOperands());
   for (Value operand : op->getOperands()) {
     // Make sure all of the operands are resolved first.
if (foldResults.empty()) {
     return markAllOverdefined(op, op->getResults());
   }
 
   // Merge the fold results into the lattice for this operation.
   assert(foldResults.size() == op->getNumResults() && "invalid result size");
   Dialect *opDialect = op->getDialect();
   for (unsigned i = 0, e = foldResults.size(); i != e; ++i) {
     LatticeValue &resultLattice = latticeValues[op->getResult(i)];
+    if (resultLattice.isOverdefined())
+      continue;
 
     // Merge in the result of the fold, either a constant or a value.
     OpFoldResult foldResult = foldResults[i];
     if (Attribute foldAttr = foldResult.dyn_cast<Attribute>())
       meet(op, resultLattice, LatticeValue(foldAttr, opDialect));
     else
       meet(op, resultLattice, latticeValues[foldResult.get<Value>()]);
   }
void SCCPSolver::visitCallOperation(CallOpInterface op) {
   auto callableLatticeIt = callableLatticeState.find(callableOp);
   if (callableLatticeIt == callableLatticeState.end())
     return markAllOverdefined(op, callResults);
 
   OperandRange callOperands = op.getArgOperands();
   auto callableArgs = callableLatticeIt->second.getCallableArguments();
   for (auto it : llvm::zip(callOperands, callableArgs)) {
     BlockArgument callableArg = std::get<1>(it);
-    if (latticeValues[callableArg].meet(latticeValues[std::get<0>(it)]))
-      visitUsers(callableArg);
+    meet(callableArg, op.getOperation()->getDialect(),
+         latticeValues[callableArg], latticeValues[std::get<0>(it)]);
   }
 
   // Merge in the lattice state for the callable results as well.
   auto callableResults = callableLatticeIt->second.getResultLatticeValues();
   for (auto it : llvm::zip(callResults, callableResults))
     meet(/*owner=*/op, /*to=*/latticeValues[std::get<0>(it)],
          /*from=*/std::get<1>(it));
 }
if (succArgs.size() != arguments.size()) {
       markAllOverdefinedAndVisitUsers(
           arguments.drop_front(firstArgIdx + succArgs.size()));
     }
 
     // Update the lattice for arguments that have inputs from the predecessor.
     OperandRange succOperands = getInputsForRegion(region->getRegionNumber());
     for (auto it : llvm::zip(succArgs, succOperands)) {
       LatticeValue &argLattice = latticeValues[std::get<0>(it)];
-      if (argLattice.meet(latticeValues[std::get<1>(it)]))
-        visitUsers(std::get<0>(it));
+      meet(std::get<0>(it), parentOp->getDialect(), argLattice,
+           latticeValues[std::get<1>(it)]);
     }
   }
 }
 
 void SCCPSolver::visitTerminatorOperation(
     Operation *op, ArrayRef<Attribute> constantOperands) {
   // If this operation has no successors, we treat it as an exiting terminator.
   if (op->getNumSuccessors() == 0) {
void SCCPSolver::markOverdefined(Value value) {
   latticeValues[value].markOverdefined();
 }
 
 bool SCCPSolver::isOverdefined(Value value) const {
   auto it = latticeValues.find(value);
   return it != latticeValues.end() && it->second.isOverdefined();
 }
 
-void SCCPSolver::meet(Operation *owner, LatticeValue &to,
+void SCCPSolver::meet(Operation *op, LatticeValue &to,
                       const LatticeValue &from) {
+  if (to.isOverdefined())
+    return;
   if (to.meet(from))
-    opWorklist.push_back(owner);
+    opWorklist.push_back(op);
+
+  // If the value resolved to a constant, merge in the dialect from op.
+  if (to.getConstant())
+    to.addConstantDialect(op->getDialect());
+}
+
+void SCCPSolver::meet(Value value, Dialect *dialect, LatticeValue &to,
+                      const LatticeValue &from) {
+  bool changed = to.meet(from);
+
+  // If the value resolved to a constant, merge in the dialect.
+  if (to.getConstant())
+    to.addConstantDialect(dialect);
+
+  // If 'to' changed, visit the users of the given value.
+  if (changed)
+    visitUsers(value);
 }
 
 //===----------------------------------------------------------------------===//
 // SCCP Pass
 //===----------------------------------------------------------------------===//
 
 namespace {
 struct SCCP : public SCCPBase<SCCP> {

mlir/test/Transforms/sccp-callgraph.mlir

func @non_symbol_defining_callable() -> i32 {
 
   %fn = "test.functional_region_op"() ({
     %1 = constant 1 : i32
     "test.return"(%1) : (i32) -> ()
   }) : () -> (() -> i32)
   %res = call_indirect %fn() : () -> (i32)
   return %res : i32
 }
+
+// -----
+
+/// Check that calls with implicit conversion have constants properly
+/// propagated.
+
+// CHECK-LABEL: func @conversion_callee
+func @conversion_callee(%arg : i32) -> (i32, i32) attributes { sym_visibility = "private" } {
+  // CHECK-DAG: %[[CONV_CST:.*]] = "test.constant"
+  // CHECK-DAG: %[[CST:.*]] = constant
+  // CHECK: return %[[CONV_CST]], %[[CST]]
+
+  %test_input = constant 0 : i32
+  return %arg, %test_input : i32, i32
+}
+
+// CHECK-LABEL: func @call_with_conversion
+func @call_with_conversion() -> (!foo.i16, !foo.i16) {
+  // CHECK: %[[CONV_CST:.*]] = "test.constant"
+  // CHECK: return %[[CONV_CST]], %[[CONV_CST]]
+
+  %test_input = "test.constant"() {value = 0 : i32} : () -> (!foo.i16)
+  %res:2 = "test.conversion_call_op"(%test_input) { callee=@conversion_callee } : (!foo.i16) -> (!foo.i16, !foo.i16)
+  return %res#0, %res#1 : !foo.i16, !foo.i16
+}

mlir/test/lib/Dialect/Test/TestDialect.cpp

 #define GET_OP_LIST
 #include "TestOps.cpp.inc"
       >();
   addInterfaces<TestOpAsmInterface, TestOpFolderDialectInterface,
                 TestInlinerInterface>();
   allowUnknownOperations();
 }
 
+Operation *TestDialect::materializeConstant(OpBuilder &builder, Attribute value,
+                                            Type type, Location loc) {
+  return builder.create<TestConstantOp>(loc, type, value);
+}
+
 LogicalResult TestDialect::verifyOperationAttribute(Operation *op,
                                                     NamedAttribute namedAttr) {
   if (namedAttr.first == "test.invalid_attr")
     return op->emitError() << "invalid to use 'test.invalid_attr'";
   return success();
 }
 
 LogicalResult TestDialect::verifyRegionArgAttribute(Operation *op,

 Optional<OperandRange> TestBranchOp::getSuccessorOperands(unsigned index) {
   assert(index == 0 && "invalid successor index");
   return getOperands();
 }
 
 bool TestBranchOp::canEraseSuccessorOperand() { return true; }
 
 //===----------------------------------------------------------------------===//
+// TestConstantOp
+//===----------------------------------------------------------------------===//
+
+OpFoldResult TestConstantOp::fold(ArrayRef<Attribute> operands) {
+  return value();
+}
+
+//===----------------------------------------------------------------------===//
 // Test IsolatedRegionOp - parse passthrough region arguments.
 //===----------------------------------------------------------------------===//
 
 static ParseResult parseIsolatedRegionOp(OpAsmParser &parser,
                                          OperationState &result) {
   OpAsmParser::OperandType argInfo;
   Type argType = parser.getBuilder().getIndexType();
 

mlir/test/lib/Dialect/Test/TestOps.td

 include "mlir/Interfaces/CallInterfaces.td"
 include "mlir/Interfaces/ControlFlowInterfaces.td"
 include "mlir/Interfaces/InferTypeOpInterface.td"
 include "mlir/Interfaces/SideEffects.td"
 
 def Test_Dialect : Dialect {
   let name = "test";
   let cppNamespace = "";
+  let hasConstantMaterializer = 1;
   let hasOperationAttrVerify = 1;
   let hasRegionArgAttrVerify = 1;
   let hasRegionResultAttrVerify = 1;
 }
 
 class TEST_Op<string mnemonic, list<OpTrait> traits = []> :
     Op<Test_Dialect, mnemonic, traits>;
 
 //===----------------------------------------------------------------------===//
+// Test Constant
+//===----------------------------------------------------------------------===//
+
+def TestConstantOp : TEST_Op<"constant", [ConstantLike]> {
+  let arguments = (ins AnyAttr:$value);
+  let results = (outs AnyType:$result);
+  let hasFolder = 1;
+}
+
+//===----------------------------------------------------------------------===//
 // Test Types
 //===----------------------------------------------------------------------===//
 
 def IntTypesOp : TEST_Op<"int_types"> {
   let results = (outs
     AnyI16:$any_i16,
     SI32:$si32,
     UI64:$ui64,