[Mlir-commits] [mlir] f151274 - [mlir][sparse] Fix Python interface for bufferization.alloc_tensor.

[llvmlistbot at llvm.org]

Author: bixia1
Date: 2022-11-09T21:38:51-08:00
New Revision: f151274454c2b55c5c64363ddc590d489d26c3fc

URL: https://github.com/llvm/llvm-project/commit/f151274454c2b55c5c64363ddc590d489d26c3fc
DIFF: https://github.com/llvm/llvm-project/commit/f151274454c2b55c5c64363ddc590d489d26c3fc.diff

LOG: [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

Added: 
    

Modified: 
    mlir/python/mlir/dialects/_bufferization_ops_ext.py
    mlir/test/Integration/Dialect/SparseTensor/taco/tools/mlir_pytaco.py

Removed: 
    


################################################################################
diff  --git a/mlir/python/mlir/dialects/_bufferization_ops_ext.py b/mlir/python/mlir/dialects/_bufferization_ops_ext.py
index 23f78fc80aec3..6ed35f4445c56 100644
--- a/mlir/python/mlir/dialects/_bufferization_ops_ext.py
+++ b/mlir/python/mlir/dialects/_bufferization_ops_ext.py
@@ -19,6 +19,7 @@ def __init__(self,
                tensor_type: Type,
                dynamic_sizes: Sequence[Value],
                copy: Value,
+               size_hint: Value,
                escape: BoolAttr,
                *,
                loc=None,
@@ -30,7 +31,7 @@ def __init__(self,
       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 b0c2bc2b9be16..f0f6ecad83f03 100644
--- 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 @@ def emit_tensor_init(self) -> ir.RankedTensorType:
     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: