[mlir] Make dispatchIndexOpFoldResult handle constant ops.

[hanchung]

Summary

This is useful when creating ops with createOrFold method. The method returns a
Value. If you store it in a SmallVector<OpFoldResult> vector, it is still a
ConstantOp if foldable, not an Attribute. Looking into the defining op will
simplify creating mixed operands cases.

Also fix a comment where uses the method.

[hanchung]

It seems like I should fix bunch of tests...

mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp

 /// In the case of a Value, a copy of the `sentinel` value is also pushed to
 /// `staticVec`. This is useful to extract mixed static and dynamic entries that
 /// come from an AttrSizedOperandSegments trait.
 static void dispatchIndexOpFoldResult(OpFoldResult ofr,
                                       SmallVectorImpl<Value> &dynamicVec,
                                       SmallVectorImpl<int64_t> &staticVec,
                                       int64_t sentinel) {
   if (auto v = ofr.dyn_cast<Value>()) {
+    if (auto cst = v.getDefiningOp<ConstantOp>()) {
+      dispatchIndexOpFoldResult(cst.getValue(), dynamicVec, staticVec,
+                                sentinel);
+      return;
+    }
     dynamicVec.push_back(v);
     staticVec.push_back(sentinel);
     return;
   }
   APInt apInt = ofr.dyn_cast<Attribute>().cast<IntegerAttr>().getValue();
   staticVec.push_back(apInt.getSExtValue());
 }
 

 //===----------------------------------------------------------------------===//
 void InitTensorOp::build(OpBuilder &b, OperationState &result,
                          ArrayRef<OpFoldResult> sizes, Type elementType,
                          ArrayRef<NamedAttribute> attrs) {
   unsigned rank = sizes.size();
   SmallVector<Value, 4> dynamicSizes;
   SmallVector<int64_t, 4> staticSizes;
   for (unsigned i = 0; i < rank; ++i) {
-    // staticLow and staticHigh have full information of the padding config.
-    // This will grow staticLow and staticHigh with 1 value. If the config is
-    // dynamic (ie not a constant), dynamicLow and dynamicHigh will grow with 1
-    // value as well.
+    // This will grow staticSizes with 1 value. If sizes[i] is dynamic (ie not a
+    // constant), dynamicSizes will grow with 1 value as well.
     dispatchIndexOpFoldResult(sizes[i], dynamicSizes, staticSizes,
                               ShapedType::kDynamicSize);
   }
-  auto resultType = RankedTensorType ::get(staticSizes, elementType);
+  auto resultType = inferResultType(staticSizes, elementType);
   build(b, result, resultType, dynamicSizes, b.getI64ArrayAttr(staticSizes));
   result.addAttributes(attrs);
 }
 
 static LogicalResult verify(InitTensorOp op) {
   RankedTensorType resultType = op.getType();
   SmallVector<int64_t, 4> staticSizes = llvm::to_vector<4>(llvm::map_range(
       op.static_sizes().cast<ArrayAttr>(),

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

 /// In the case of a Value, a copy of the `sentinel` value is also pushed to
 /// `staticVec`. This is useful to extract mixed static and dynamic entries that
 /// come from an AttrSizedOperandSegments trait.
 static void dispatchIndexOpFoldResult(OpFoldResult ofr,
                                       SmallVectorImpl<Value> &dynamicVec,
                                       SmallVectorImpl<int64_t> &staticVec,
                                       int64_t sentinel) {
   if (auto v = ofr.dyn_cast<Value>()) {
+    if (auto cst = v.getDefiningOp<ConstantOp>()) {
+      dispatchIndexOpFoldResult(cst.getValue(), dynamicVec, staticVec,
+                                sentinel);
+      return;
+    }
     dynamicVec.push_back(v);
     staticVec.push_back(sentinel);
     return;
   }
   APInt apInt = ofr.dyn_cast<Attribute>().cast<IntegerAttr>().getValue();
   staticVec.push_back(apInt.getSExtValue());
 }