[MLIR] Add `OpFoldResult` builder to tensor.dim op

[pifon2a]

Why is this needed? DimOp always takes Value for the dimension index. There is no index attribute, which would allow to fold constants into the dimOp. With the current DimOp it is easier to call getValueOrCreateConstantIndexOp when DimOp::build is called.

[frgossen]

Right, getValueOrCreateConstantIndexOp works.
As there is a builder for int64_t, which creates a constant one for OpFoldResult would be the natural completion, no?

[frgossen]

Don't feel strong about this though, so we can discard this also.

mlir/include/mlir/Dialect/Tensor/IR/TensorOps.td

let arguments = (ins AnyTensor:$source,
                        Index:$index);
   let results = (outs Index:$result);
 
   let assemblyFormat = [{
     attr-dict $source `,` $index `:` type($source)
   }];
 
   let builders = [
-    OpBuilder<(ins "Value":$source, "int64_t":$index)>
+    OpBuilder<(ins "Value":$source, "int64_t":$index)>,
+    OpBuilder<(ins "Value":$source, "OpFoldResult":$index)>,
   ];
 
   let extraClassDeclaration = [{
     /// Helper function to get the index as a simple integer if it is constant.
     std::optional<int64_t> getConstantIndex();
 
     /// Interface method of ShapedDimOpInterface: Return the source tensor.
     Value getShapedValue() { return getSource(); }

mlir/lib/Dialect/Tensor/IR/TensorOps.cpp

 
 void DimOp::build(OpBuilder &builder, OperationState &result, Value source,
                   int64_t index) {
   auto loc = result.location;
   Value indexValue = builder.create<arith::ConstantIndexOp>(loc, index);
   build(builder, result, source, indexValue);
 }
 
+void DimOp::build(OpBuilder &builder, OperationState &result, Value source,
+                  OpFoldResult index) {
+  if (index.is<Value>())
+    return build(builder, result, source, index.get<Value>());
+  auto indexAttr = index.get<Attribute>();
+  return build(builder, result, source, indexAttr.cast<IntegerAttr>().getInt());
+}
+
 std::optional<int64_t> DimOp::getConstantIndex() {
   return getConstantIntValue(getIndex());
 }
 
 Speculation::Speculatability DimOp::getSpeculatability() {
   auto constantIndex = getConstantIndex();
   if (!constantIndex)
     return Speculation::NotSpeculatable;