[mlir][sparse] Fix Python interface for bufferization.alloc_tensor.

Add size_hint operand to the Python interface.
Fix pytaco.

Reviewed By: aartbik

Differential Revision: https://reviews.llvm.org/D137754

diff --git a/mlir/python/mlir/dialects/_bufferization_ops_ext.py b/mlir/python/mlir/dialects/_bufferization_ops_ext.py
index 23f78fc..6ed35f4 100644 (file)
--- a/mlir/python/mlir/dialects/_bufferization_ops_ext.py
+++ b/mlir/python/mlir/dialects/_bufferization_ops_ext.py
@@ -19,6 +19,7 @@ class AllocTensorOp:
                tensor_type: Type,
                dynamic_sizes: Sequence[Value],
                copy: Value,
+               size_hint: Value,
                escape: BoolAttr,
                *,
                loc=None,
@@ -30,7 +31,7 @@ class AllocTensorOp:
       attributes["escape"] = escape
     op = self.build_generic(
         results=[tensor_type],
-        operands=[dynamic_sizes, copy],
+        operands=[dynamic_sizes, copy, size_hint],
         attributes=attributes,
         loc=loc,
         ip=ip)

diff --git a/mlir/test/Integration/Dialect/SparseTensor/taco/tools/mlir_pytaco.py b/mlir/test/Integration/Dialect/SparseTensor/taco/tools/mlir_pytaco.py
index b0c2bc2..f0f6eca 100644 (file)
--- a/mlir/test/Integration/Dialect/SparseTensor/taco/tools/mlir_pytaco.py
+++ b/mlir/test/Integration/Dialect/SparseTensor/taco/tools/mlir_pytaco.py
@@ -909,7 +909,7 @@ class _StructOpInfo:
     mlir_type = _mlir_tensor_type(self.dst_dtype, self.dst_dims,
                                   self.dst_format.mlir_tensor_attr())
     index_type = ir.IndexType.get()
-    return bufferization.AllocTensorOp(mlir_type, [], None, None)
+    return bufferization.AllocTensorOp(mlir_type, [], None, None, None)
 
 
 class _Stats: