[mlir] Add canonicalization for extract(tensor.from_elements) in 0d case.

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

struct ExtractElementFromTensorFromElements
 
   LogicalResult matchAndRewrite(tensor::ExtractOp extract,
                                 PatternRewriter &rewriter) const final {
     auto tensorFromElements = extract.tensor().getDefiningOp<FromElementsOp>();
     if (!tensorFromElements)
       return failure();
     auto tensorType = tensorFromElements.getType().cast<RankedTensorType>();
     auto rank = tensorType.getRank();
+    if (rank == 0) {
+      rewriter.replaceOp(extract, tensorFromElements.getOperand(0));
+      return success();
+    }
     SmallVector<APInt, 3> indices(rank);
     int64_t flatIndex = 0;
     int64_t stride = 1;
     for (int i = rank - 1; i >= 0; --i) {
       APInt index;
       if (!matchPattern(extract.indices()[i], m_ConstantInt(&index)))
         return failure();
       if (i < rank - 1)

mlir/test/Dialect/Tensor/canonicalize.mlir

func @extract_from_tensor.from_elements(%element : index) -> index {
   %tensor = tensor.from_elements %element : tensor<1xindex>
   %extracted_element = tensor.extract %tensor[%c0] : tensor<1xindex>
   // CHECK: [[ARG]] : index
   return %extracted_element : index
 }
 
 // -----
 
+// CHECK-LABEL: func @extract_from_tensor.from_elements_0d
+func @extract_from_tensor.from_elements_0d(%element : index) -> index {
+  // CHECK-SAME: ([[ARG:%.*]]: index)
+  %c0 = arith.constant 0 : index
+  %tensor = tensor.from_elements %element : tensor<index>
+  %extracted_element = tensor.extract %tensor[] : tensor<index>
+  // CHECK: [[ARG]] : index
+  return %extracted_element : index
+}
+
+// -----
+
 // CHECK-LABEL: func @extract_from_tensor.from_elements_3d
 func @extract_from_tensor.from_elements_3d()
     -> (f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32) {
   %f0 = arith.constant 0.0 : f32
   %f1 = arith.constant 1.0 : f32
   %f2 = arith.constant 2.0 : f32
   %f3 = arith.constant 3.0 : f32
   %f4 = arith.constant 4.0 : f32