propagate poison as part of the interface

Author: Hardcode84

mlir/include/mlir/Dialect/Arith/IR/ArithOps.td

@@ -1687,7 +1687,7 @@ class BooleanConditionOrMatchingShape<string condition, string result> :
 def SelectOp : Arith_Op<"select", [Pure,
     AllTypesMatch<["true_value", "false_value", "result"]>,
     BooleanConditionOrMatchingShape<"condition", "result">,
-    DeclareOpInterfaceMethods<InferIntRangeInterface, ["inferResultRangesFromOptional"]>,
+    DeclareOpInterfaceMethods<InferIntRangeInterface, ["inferResultRangesOrPoison"]>,
     DeclareOpInterfaceMethods<SelectLikeOpInterface>]> {
   let summary = "select operation";
   let description = [{

mlir/include/mlir/Dialect/Vector/IR/VectorOps.td

@@ -843,7 +843,7 @@ def Vector_FromElementsOp : Vector_Op<"from_elements", [
 
 def Vector_InsertOp :
   Vector_Op<"insert", [Pure,
-     DeclareOpInterfaceMethods<InferIntRangeInterface, ["inferResultRanges"]>,
+     DeclareOpInterfaceMethods<InferIntRangeInterface, ["inferResultRangesOrPoison"]>,
      PredOpTrait<"source operand and result have same element type",
                  TCresVTEtIsSameAsOpBase<0, 0>>,
      AllTypesMatch<["dest", "result"]>]> {

mlir/include/mlir/IR/Matchers.h

@@ -129,7 +129,7 @@ struct infer_int_range_op_binder {
       *bind_value = argRanges;
       matched = true;
     };
-    inferIntRangeOp.inferResultRangesFromOptional(argRanges, setResultRanges);
+    inferIntRangeOp.inferResultRangesOrPoison(argRanges, setResultRanges);
     return matched;
   }
 };

mlir/include/mlir/Interfaces/InferIntRangeInterface.h

@@ -198,6 +198,13 @@ void defaultInferResultRanges(InferIntRangeInterface interface,
 void defaultInferResultRangesFromOptional(InferIntRangeInterface interface,
                                           ArrayRef<ConstantIntRanges> argRanges,
                                           SetIntRangeFn setResultRanges);
+
+/// Default implementation of `inferResultRangesOrPoison` which propagates
+/// poison and dispatches to the `inferResultRangesFromOptional`.
+void defaultInferResultRangesOrPoison(InferIntRangeInterface interface,
+                                      ArrayRef<IntegerValueRange> argRanges,
+                                      SetIntLatticeFn setResultRanges);
+
 } // end namespace intrange::detail
 } // end namespace mlir
 

mlir/include/mlir/Interfaces/InferIntRangeInterface.td

@@ -33,6 +33,9 @@ def InferIntRangeInterface : OpInterface<"InferIntRangeInterface"> {
       When operations take non-integer inputs, the
      `inferResultRangesFromOptional` method should be implemented instead.
 
+      If any of the operands have poison ranges, they will be propagated to the
+      results automatically after the metdod returns.
+
       When called on an op that also implements the RegionBranchOpInterface
       or BranchOpInterface, this method should not attempt to infer the values
       of the branch results, as this will be handled by the analyses that use
@@ -60,6 +63,9 @@ def InferIntRangeInterface : OpInterface<"InferIntRangeInterface"> {
       as an argument. When implemented, `setValueRange` should be called on
       all result values for the operation.
 
+      If any of the operands have poison ranges, they will be propagated to the
+      results automatically after the metdod returns.
+
       This method allows for more precise implementations when operations
       want to reason about inputs which may be undefined during the analysis.
     }],
@@ -72,6 +78,30 @@ def InferIntRangeInterface : OpInterface<"InferIntRangeInterface"> {
       ::mlir::intrange::detail::defaultInferResultRanges($_op,
                                                          argRanges,
                                                          setResultRanges);
+    }]>,
+
+    InterfaceMethod<[{
+      Infer the bounds on the results of this op given the lattice representation
+      of the bounds for its arguments. For each result value or block argument
+      (that isn't a branch argument, since the dataflow analysis handles
+      those case), the method should call `setValueRange` with that `Value`
+      as an argument. When implemented, `setValueRange` should be called on
+      all result values for the operation.
+
+      Unlike `inferResultRanges`/`inferResultRangesFromOptional` this method
+      does not automatically propagate poison from the inputs. This allows more
+      precise poison semantics implementation.
+    }],
+    /*retTy=*/"void",
+    /*methodName=*/"inferResultRangesOrPoison",
+    /*args=*/(ins "::llvm::ArrayRef<::mlir::IntegerValueRange>":$argRanges,
+                  "::mlir::SetIntLatticeFn":$setResultRanges),
+    /*methodBody=*/"",
+    /*defaultImplementation=*/[{
+      ::mlir::intrange::detail::defaultInferResultRangesOrPoison(
+        $_op,
+        argRanges,
+        setResultRanges);
     }]>
   ];
 }

mlir/lib/Analysis/DataFlow/IntegerRangeAnalysis.cpp

@@ -124,7 +124,7 @@ LogicalResult IntegerRangeAnalysis::visitOperation(
     propagateIfChanged(lattice, changed);
   };
 
-  inferrable.inferResultRangesFromOptional(argRanges, joinCallback);
+  inferrable.inferResultRangesOrPoison(argRanges, joinCallback);
   return success();
 }
 
@@ -167,7 +167,7 @@ void IntegerRangeAnalysis::visitNonControlFlowArguments(
       propagateIfChanged(lattice, changed);
     };
 
-    inferrable.inferResultRangesFromOptional(argRanges, joinCallback);
+    inferrable.inferResultRangesOrPoison(argRanges, joinCallback);
     return;
   }
 

mlir/lib/Dialect/Arith/IR/InferIntRangeInterfaceImpls.cpp

@@ -312,7 +312,7 @@ void arith::CmpIOp::inferResultRanges(ArrayRef<ConstantIntRanges> argRanges,
 // SelectOp
 //===----------------------------------------------------------------------===//
 
-void arith::SelectOp::inferResultRangesFromOptional(
+void arith::SelectOp::inferResultRangesOrPoison(
     ArrayRef<IntegerValueRange> argRanges, SetIntLatticeFn setResultRange) {
   std::optional<APInt> mbCondVal =
       argRanges[0].isUninitialized()

mlir/lib/Dialect/Vector/IR/VectorOps.cpp

@@ -3207,9 +3207,14 @@ void ShuffleOp::getCanonicalizationPatterns(RewritePatternSet &results,
 // InsertOp
 //===----------------------------------------------------------------------===//
 
-void vector::InsertOp::inferResultRanges(ArrayRef<ConstantIntRanges> argRanges,
-                                         SetIntRangeFn setResultRanges) {
-  setResultRanges(getResult(), argRanges[0].rangeUnion(argRanges[1]));
+void vector::InsertOp::inferResultRangesOrPoison(
+    ArrayRef<IntegerValueRange> argRanges, SetIntLatticeFn setResultRanges) {
+  if (argRanges[0].isUninitialized() || argRanges[1].isUninitialized())
+    return;
+
+  const ConstantIntRanges &range0 = argRanges[0].getValue();
+  const ConstantIntRanges &range1 = argRanges[1].getValue();
+  setResultRanges(getResult(), range0.rangeUnion(range1));
 }
 
 void vector::InsertOp::build(OpBuilder &builder, OperationState &result,

mlir/lib/Interfaces/InferIntRangeInterface.cpp

@@ -253,3 +253,38 @@ void mlir::intrange::detail::defaultInferResultRangesFromOptional(
           setResultRanges(value, argRanges.getValue());
       });
 }
+
+void mlir::intrange::detail::defaultInferResultRangesOrPoison(
+    InferIntRangeInterface interface, ArrayRef<IntegerValueRange> argRanges,
+    SetIntLatticeFn setResultRanges) {
+
+  bool signedPoison = false;
+  bool unsignedPoison = false;
+  for (const IntegerValueRange &range : argRanges) {
+    if (range.isUninitialized())
+      continue;
+
+    const ConstantIntRanges &value = range.getValue();
+    signedPoison = signedPoison || value.isSignedPoison();
+    unsignedPoison = unsignedPoison || value.isUnsignedPoison();
+  }
+
+  auto visitor = [&](Value value, const IntegerValueRange &range) {
+    if (range.isUninitialized())
+      return;
+
+    if (!signedPoison && !unsignedPoison)
+      return setResultRanges(value, range);
+
+    const ConstantIntRanges &origRange = range.getValue();
+    auto poison = ConstantIntRanges::poison(origRange.getBitWidth());
+    APInt umin = unsignedPoison ? poison.umin() : origRange.umin();
+    APInt umax = unsignedPoison ? poison.umax() : origRange.umax();
+    APInt smin = signedPoison ? poison.smin() : origRange.smin();
+    APInt smax = signedPoison ? poison.smax() : origRange.smax();
+
+    setResultRanges(value, ConstantIntRanges(umin, umax, smin, smax));
+  };
+
+  interface.inferResultRangesFromOptional(argRanges, visitor);
+}