From 0a02f76d11d727ca89a1cd3d29e0b867de7051d9 Mon Sep 17 00:00:00 2001
From: max <maksim.levental@gmail.com>
Date: Sun, 30 Apr 2023 17:46:29 -0500
Subject: [PATCH] [MLIR][tensor] generate default builder for FromElementsOp

Removed builder is the same as default builder, with the added benefit that python bindings will be generated for the default builder.

Reviewed By: ftynse

Differential Revision: https://reviews.llvm.org/D149508
---
 mlir/include/mlir/Dialect/Tensor/IR/TensorOps.td |  2 --
 mlir/lib/Dialect/Tensor/IR/TensorOps.cpp         |  6 -----
 mlir/test/python/dialects/tensor.py              | 32 +++++++++++++++++++++---
 3 files changed, 29 insertions(+), 11 deletions(-)

diff --git a/mlir/include/mlir/Dialect/Tensor/IR/TensorOps.td b/mlir/include/mlir/Dialect/Tensor/IR/TensorOps.td
index 0c589ce..db979c0 100644
--- a/mlir/include/mlir/Dialect/Tensor/IR/TensorOps.td
+++ b/mlir/include/mlir/Dialect/Tensor/IR/TensorOps.td
@@ -484,9 +484,7 @@ def Tensor_FromElementsOp : Tensor_Op<"from_elements", [
 
   let assemblyFormat = "$elements attr-dict `:` type($result)";
 
-  let skipDefaultBuilders = 1;
   let builders = [
-    OpBuilder<(ins "Type":$resultType, "ValueRange":$elements)>,
     // Special case builder for when `elements` has size >=1.
     OpBuilder<(ins "ValueRange":$elements)>
   ];
diff --git a/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp b/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp
index 99382a3..33639e1 100644
--- a/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp
+++ b/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp
@@ -853,12 +853,6 @@ void FromElementsOp::getAsmResultNames(
 }
 
 void FromElementsOp::build(OpBuilder &builder, OperationState &result,
-                           Type resultType, ValueRange elements) {
-  result.addOperands(elements);
-  result.addTypes(resultType);
-}
-
-void FromElementsOp::build(OpBuilder &builder, OperationState &result,
                            ValueRange elements) {
   assert(!elements.empty() && "expected at least one element");
   Type resultType = RankedTensorType::get(
diff --git a/mlir/test/python/dialects/tensor.py b/mlir/test/python/dialects/tensor.py
index d8ea426..6c85dd1 100644
--- a/mlir/test/python/dialects/tensor.py
+++ b/mlir/test/python/dialects/tensor.py
@@ -82,7 +82,6 @@ def testInferTypesInsertSlice():
   with Context() as ctx, Location.unknown():
     module = Module.create()
     f32Type = F32Type.get()
-    indexType = IndexType.get()
     with InsertionPoint(module.body):
 
       @func.FuncOp.from_py_func(
@@ -92,8 +91,6 @@ def testInferTypesInsertSlice():
       # CHECK:      tensor.insert_slice %arg0 into %arg1[0, 0] [1, 1] [0, 0] :
       # CHECK-SAME:   tensor<1x1xf32> into tensor<1x1xf32>
       def f(source, dest):
-        c0 = arith.ConstantOp(indexType, 0)
-        c1 = arith.ConstantOp(indexType, 1)
         d0 = tensor.InsertSliceOp(source, dest, [], [], [],
                                   DenseI64ArrayAttr.get([0, 0]),
                                   DenseI64ArrayAttr.get([1, 1]),
@@ -101,3 +98,32 @@ def testInferTypesInsertSlice():
         return [d0.result]
 
   print(module)
+
+
+# CHECK-LABEL: TEST: testFromElementsOp
+@run
+def testFromElementsOp():
+  with Context() as ctx, Location.unknown():
+    module = Module.create()
+    f32 = F32Type.get()
+    with InsertionPoint(module.body):
+      @func.FuncOp.from_py_func()
+      def default_builder():
+        c0 = arith.ConstantOp(f32, 0.0)
+        # CHECK: %[[C0:.*]] = "arith.constant"() {value = 0.000000e+00 : f32} : () -> f32
+        print(c0)
+        c1 = arith.ConstantOp(f32, 1.0)
+        # CHECK: %[[C1:.*]] = "arith.constant"() {value = 1.000000e+00 : f32} : () -> f32
+        print(c1)
+
+        t = tensor.FromElementsOp(RankedTensorType.get((2,), f32), [c0, c1])
+        # CHECK: %{{.*}} = "tensor.from_elements"(%[[C0]], %[[C1]]) : (f32, f32) -> tensor<2xf32>
+        print(t)
+
+        t = tensor.FromElementsOp(RankedTensorType.get((2, 1), f32), [c0, c1])
+        # CHECK: %{{.*}} = "tensor.from_elements"(%[[C0]], %[[C1]]) : (f32, f32) -> tensor<2x1xf32>
+        print(t)
+
+        t = tensor.FromElementsOp(RankedTensorType.get((1, 2), f32), [c0, c1])
+        # CHECK: %{{.*}} = "tensor.from_elements"(%[[C0]], %[[C1]]) : (f32, f32) -> tensor<1x2xf32>
+        print(t)
-- 
2.7.4