[mlir][Tensor] Add folding for tensor.from_elements

[bkramer]

Summary

This trivially folds into a constant when all operands are constant.

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

def Tensor_FromElementsOp : Tensor_Op<"from_elements", [
   let skipDefaultBuilders = 1;
   let builders = [
     OpBuilder<(ins "Type":$elementType, "ValueRange":$elements)>,
     // Special case builder for when `elements` has size >=1.
     OpBuilder<(ins "ValueRange":$elements)>
   ];
 
   let hasCanonicalizer = 1;
+  let hasFolder = 1;
 }
 
 //===----------------------------------------------------------------------===//
 // GenerateOp
 //===----------------------------------------------------------------------===//
 
 def Tensor_GenerateOp : Tensor_Op<"generate",
     [RecursiveSideEffects,

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

 }
 
 void FromElementsOp::build(OpBuilder &builder, OperationState &result,
                            ValueRange elements) {
   assert(!elements.empty() && "expected at least one element");
   build(builder, result, elements.front().getType(), elements);
 }
 
+OpFoldResult FromElementsOp::fold(ArrayRef<Attribute> operands) {
+  if (!llvm::is_contained(operands, nullptr))
+    return DenseElementsAttr::get(getType(), operands);
+  return {};
+}
+
 namespace {
 
 // Canonicalizes the pattern of the form
 //
 // %tensor = tensor.from_elements(%element) : (i32) -> tensor<1xi32>
 // %extracted_element = tensor.extract %tensor[%c0] : tensor<1xi32>
 //
 // to just %element.

mlir/test/Dialect/Linalg/detensorize_trivial.mlir

 // DET-ALL-NEXT:    tensor.extract %{{.*}}[]
 // DET-ALL-NEXT:    cmpi slt, %{{.*}}, %{{.*}}
 // DET-ALL-NEXT:    tensor.from_elements %{{.*}}
 // DET-ALL-NEXT:    linalg.tensor_reshape %{{.*}}
 // DET-ALL-NEXT:    return %{{.*}} : tensor<i1>
 // DET-ALL-NEXT:  }
 
 // DET-CF-LABEL: func @main(%{{.*}}: tensor<i32>)
-// DET-CF-NEXT:    constant 10 : i32
-// DET-CF-NEXT:    tensor.from_elements %{{.*}}
-// DET-CF-NEXT:    linalg.tensor_reshape %{{.*}}
+// DET-CF-NEXT:    constant dense<10> : tensor<i32>
 // DET-CF-NEXT:    linalg.init_tensor [] : tensor<i1>
 // DET-CF-NEXT:    linalg.generic
 // DET-CF-NEXT:    ^{{.*}}(%{{.*}}: i32, %{{.*}}: i32, %{{.*}}: i1)
 // DET-CF-NEXT:      cmpi slt, %{{.*}}, %{{.*}}
 // DET-CF-NEXT:      linalg.yield %{{.*}}
 // DET-CF-NEXT:    } -> tensor<i1>
 // DET-CF-NEXT:    return %{{.*}}
 // DET-CF-NEXT:  }

mlir/test/Dialect/Tensor/canonicalize.mlir

%0 = tensor.generate %size1, %c5, %size4 {
     ^bb0(%arg0: index, %arg1: index, %arg2: index, %arg3: index, %arg4: index):
     %1 = constant 32 : index
     tensor.yield %1 : index
   // CHECK: : tensor<3x?x5x7x?xindex>
   } : tensor<3x?x?x7x?xindex>
   // CHECK: tensor.cast %{{.*}} : tensor<3x?x5x7x?xindex> to tensor<3x?x?x7x?xindex>
   return %0 : tensor<3x?x?x7x?xindex>
 }
+
+// -----
+
+// CHECK-LABEL: @from_elements.constant
+func @from_elements.constant() -> tensor<3xindex> {
+  // CHECK: %[[CST:.*]] = constant dense<[1, 2, 1]> : tensor<3xindex>
+  // CHECK: return %[[CST]]
+  %c1 = constant 1 : index
+  %c2 = constant 2 : index
+  %tensor = tensor.from_elements %c1, %c2, %c1 : tensor<3xindex>
+  return %tensor : tensor<3xindex>
+}