diff --git a/mlir/include/mlir/Dialect/Bufferization/IR/BufferizationOps.td b/mlir/include/mlir/Dialect/Bufferization/IR/BufferizationOps.td --- a/mlir/include/mlir/Dialect/Bufferization/IR/BufferizationOps.td +++ b/mlir/include/mlir/Dialect/Bufferization/IR/BufferizationOps.td @@ -29,13 +29,30 @@ let summary = "buffer allocation in tensor land"; let description = [{ - `bufferization.alloc_tensor` is an operation that bufferizes to a buffer - allocation of a given shape. The shape could be dynamic or static. - Reading from the result of an `alloc_tensor` op yields an undefined value. + `bufferization.alloc_tensor` materializes an uninitialized tensor with a + given shape (dynamic or static). It always bufferizes to a new buffer + allocation of the given shape. Reading from the result of an `alloc_tensor` + op yields an undefined value. + + `alloc_tensor` is a helper op for bufferization. The operation is provided + as an anchor that marks the beginning of a new tensor SSA use-def chain. It + can be used to control in-place bufferization decisions during One-Shot + Bufferize: The bufferized result of a `bufferization.alloc_tensor` does not + alias with any other buffer, so it can be used to resolve read-after-write + conflicts that would have been introduced by the in-place bufferization of + another op. + + Both dense and sparse tensor types are supported. The result of a + `bufferization.alloc_tensor` is a tensor value that can be used like any + other tensor value. In practice, it is often used as the "out" operand of + another op. E.g.: - `alloc_tensor` is a helper op for bufferization. It marks the beginning of - a new tensor SSA use-def chain and is used to control in-place bufferization - decisions during One-Shot Bufferize. + ```mlir + %c = sparse_tensor.init_tensor [%d1, %d2] : tensor + %0 = linalg.matmul + ins(%a, %b: tensor, tensor) + outs(%c: tensor) -> tensor + ``` }]; let arguments = (ins Variadic:$dynamicSizes); diff --git a/mlir/include/mlir/Dialect/SparseTensor/IR/SparseTensorOps.td b/mlir/include/mlir/Dialect/SparseTensor/IR/SparseTensorOps.td --- a/mlir/include/mlir/Dialect/SparseTensor/IR/SparseTensorOps.td +++ b/mlir/include/mlir/Dialect/SparseTensor/IR/SparseTensorOps.td @@ -49,29 +49,6 @@ let hasVerifier = 1; } -def SparseTensor_InitOp : SparseTensor_Op<"init", [NoSideEffect]>, - Arguments<(ins Variadic:$sizes)>, - Results<(outs AnyTensor:$result)> { - string summary = "Materializes an unitialized sparse tensor"; - string description = [{ - Materializes an uninitialized sparse tensor with given shape (either static - or dynamic). The operation is provided as an anchor that materializes a - properly typed but uninitialized sparse tensor into the output clause of - a subsequent operation that yields a sparse tensor as the result. - - Example: - - ```mlir - %c = sparse_tensor.init_tensor [%d1, %d2] : tensor - %0 = linalg.matmul - ins(%a, %b: tensor, tensor) - outs(%c: tensor) -> tensor - ``` - }]; - let assemblyFormat = "`[` $sizes `]` attr-dict `:` type($result)"; - let hasVerifier = 1; -} - def SparseTensor_ConvertOp : SparseTensor_Op<"convert", [NoSideEffect, SameOperandsAndResultElementType]>, Arguments<(ins AnyTensor:$source)>, @@ -417,7 +394,7 @@ Example of isEqual applied to intersecting elements only. ```mlir - %C = sparse_tensor.init... + %C = bufferization.alloc_tensor... %0 = linalg.generic #trait ins(%A: tensor, %B: tensor) outs(%C: tensor) { @@ -438,7 +415,7 @@ Example of A+B in upper triangle, A-B in lower triangle (not working yet, but construct will be available soon). ```mlir - %C = sparse_tensor.init... + %C = bufferization.alloc_tensor... %1 = linalg.generic #trait ins(%A: tensor, %B: tensor outs(%C: tensor { @@ -470,7 +447,7 @@ is *not* overlapped by B. The element type of B can be different than A because we never use its values, only its sparse structure. ```mlir - %C = sparse_tensor.init... + %C = bufferization.alloc_tensor... %2 = linalg.generic #trait ins(%A: tensor, %B: tensor outs(%C: tensor { @@ -517,7 +494,7 @@ Example of A+1, restricted to existing elements: ```mlir - %C = sparse_tensor.init... + %C = bufferization.alloc_tensor... %0 = linalg.generic #trait ins(%A: tensor) outs(%C: tensor) { diff --git a/mlir/lib/Dialect/Linalg/Transforms/SparseTensorRewriting.cpp b/mlir/lib/Dialect/Linalg/Transforms/SparseTensorRewriting.cpp --- a/mlir/lib/Dialect/Linalg/Transforms/SparseTensorRewriting.cpp +++ b/mlir/lib/Dialect/Linalg/Transforms/SparseTensorRewriting.cpp @@ -16,6 +16,7 @@ //===----------------------------------------------------------------------===// #include "mlir/Dialect/Arithmetic/IR/Arithmetic.h" +#include "mlir/Dialect/Bufferization/IR/Bufferization.h" #include "mlir/Dialect/Linalg/IR/Linalg.h" #include "mlir/Dialect/Linalg/Transforms/Transforms.h" #include "mlir/Dialect/SparseTensor/IR/SparseTensor.h" @@ -24,6 +25,7 @@ #include "mlir/Support/LLVM.h" using namespace mlir; +using namespace mlir::bufferization; using namespace mlir::linalg; using namespace mlir::sparse_tensor; @@ -47,7 +49,8 @@ static bool isEmptyInit(OpOperand *op) { Value val = op->get(); return matchPattern(val, m_Zero()) || matchPattern(val, m_AnyZeroFloat()) || - val.getDefiningOp() || val.getDefiningOp(); + val.getDefiningOp() || + val.getDefiningOp(); } // Helper to detect sampling operation. diff --git a/mlir/lib/Dialect/SparseTensor/IR/SparseTensorDialect.cpp b/mlir/lib/Dialect/SparseTensor/IR/SparseTensorDialect.cpp --- a/mlir/lib/Dialect/SparseTensor/IR/SparseTensorDialect.cpp +++ b/mlir/lib/Dialect/SparseTensor/IR/SparseTensorDialect.cpp @@ -214,28 +214,6 @@ return success(); } -LogicalResult InitOp::verify() { - if (!getSparseTensorEncoding(result().getType())) - return emitError("expected a sparse tensor result"); - RankedTensorType ttp = getType().cast(); - unsigned rank = ttp.getRank(); - if (rank != sizes().size()) - return emitError("unexpected mismatch between tensor rank and sizes: ") - << rank << " vs. " << sizes().size(); - auto shape = ttp.getShape(); - for (unsigned i = 0; i < rank; i++) { - if (shape[i] == ShapedType::kDynamicSize) - continue; - IntegerAttr constantAttr; - if (!matchPattern(sizes()[i], m_Constant(&constantAttr)) || - constantAttr.getInt() != shape[i]) { - return emitError("unexpected mismatch with static dimension size ") - << shape[i]; - } - } - return success(); -} - LogicalResult ConvertOp::verify() { if (auto tp1 = source().getType().dyn_cast()) { if (auto tp2 = dest().getType().dyn_cast()) { diff --git a/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorConversion.cpp b/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorConversion.cpp --- a/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorConversion.cpp +++ b/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorConversion.cpp @@ -459,22 +459,33 @@ } }; -/// Sparse conversion rule for the init operator. -class SparseTensorInitConverter : public OpConversionPattern { +/// Sparse conversion rule for the alloc operator. +class SparseTensorAllocConverter + : public OpConversionPattern { using OpConversionPattern::OpConversionPattern; LogicalResult - matchAndRewrite(InitOp op, OpAdaptor adaptor, + matchAndRewrite(bufferization::AllocTensorOp op, OpAdaptor adaptor, ConversionPatternRewriter &rewriter) const override { - Type resType = op.getType(); + RankedTensorType resType = op.getType(); auto enc = getSparseTensorEncoding(resType); if (!enc) return failure(); + // Gather all dimension sizes as SSA values. + SmallVector sizes; + unsigned int operandCtr = 0; + for (int64_t i = 0; i < resType.getRank(); ++i) { + if (resType.isDynamicDim(i)) { + sizes.push_back(adaptor.getOperands()[operandCtr++]); + } else { + sizes.push_back(rewriter.create( + op.getLoc(), op.getStaticSize(i))); + } + } // Generate the call to construct empty tensor. The sizes are - // explicitly defined by the arguments to the init operator. + // explicitly defined by the arguments to the alloc operator. SmallVector params; ShapedType stp = resType.cast(); - newParams(rewriter, params, op, stp, enc, Action::kEmpty, - adaptor.getOperands()); + newParams(rewriter, params, op, stp, enc, Action::kEmpty, sizes); rewriter.replaceOp(op, genNewCall(rewriter, op, params)); return success(); } @@ -912,7 +923,7 @@ const SparseTensorConversionOptions &options) { patterns.add(); + target.addIllegalOp(); // Translate strategy flags to strategy options. SparseTensorConversionOptions options( sparseToSparseConversionStrategy(sparseToSparse)); diff --git a/mlir/lib/Dialect/SparseTensor/Transforms/Sparsification.cpp b/mlir/lib/Dialect/SparseTensor/Transforms/Sparsification.cpp --- a/mlir/lib/Dialect/SparseTensor/Transforms/Sparsification.cpp +++ b/mlir/lib/Dialect/SparseTensor/Transforms/Sparsification.cpp @@ -293,7 +293,7 @@ /// Returns true if tensor materializes uninitialized into the computation. static bool isMaterializing(Value val) { return val.getDefiningOp() || - val.getDefiningOp(); + val.getDefiningOp(); } /// Returns true when the tensor expression is admissable for codegen. diff --git a/mlir/test/Dialect/SparseTensor/conversion.mlir b/mlir/test/Dialect/SparseTensor/conversion.mlir --- a/mlir/test/Dialect/SparseTensor/conversion.mlir +++ b/mlir/test/Dialect/SparseTensor/conversion.mlir @@ -135,7 +135,7 @@ // CHECK: %[[T:.*]] = call @newSparseTensor(%[[X]], %[[Y]], %[[Z]], %{{.*}}, %{{.*}}, %{{.*}}, %[[Empty]], %[[NP]]) // CHECK: return %[[T]] : !llvm.ptr func.func @sparse_init(%arg0: index, %arg1: index) -> tensor { - %0 = sparse_tensor.init [%arg0, %arg1] : tensor + %0 = bufferization.alloc_tensor(%arg0, %arg1) : tensor return %0 : tensor } @@ -512,8 +512,7 @@ // CHECK-DAG: linalg.fill ins(%{{.*}} : i1) outs(%[[B]] : memref) // CHECK: return %[[C]] : memref func.func @sparse_expansion() -> memref { - %c = arith.constant 8 : index - %0 = sparse_tensor.init [%c, %c] : tensor<8x8xf64, #SparseMatrix> + %0 = bufferization.alloc_tensor() : tensor<8x8xf64, #SparseMatrix> %values, %filled, %added, %count = sparse_tensor.expand %0 : tensor<8x8xf64, #SparseMatrix> to memref, memref, memref, index return %added : memref diff --git a/mlir/test/Dialect/SparseTensor/invalid.mlir b/mlir/test/Dialect/SparseTensor/invalid.mlir --- a/mlir/test/Dialect/SparseTensor/invalid.mlir +++ b/mlir/test/Dialect/SparseTensor/invalid.mlir @@ -16,36 +16,6 @@ // ----- -func.func @invalid_init_dense(%arg0: index, %arg1: index) -> tensor { - // expected-error@+1 {{expected a sparse tensor result}} - %0 = sparse_tensor.init [%arg0, %arg1] : tensor - return %0 : tensor -} - -// ----- - -#SparseVector = #sparse_tensor.encoding<{dimLevelType = ["compressed"]}> - -func.func @invalid_init_rank(%arg0: index) -> tensor { - // expected-error@+1 {{unexpected mismatch between tensor rank and sizes: 1 vs. 2}} - %0 = sparse_tensor.init [%arg0, %arg0] : tensor - return %0 : tensor -} - -// ----- - -#SparseMatrix = #sparse_tensor.encoding<{dimLevelType = ["compressed", "compressed"]}> - -func.func @invalid_init_size() -> tensor { - %c10 = arith.constant 10 : index - %c20 = arith.constant 20 : index - // expected-error@+1 {{unexpected mismatch with static dimension size 10}} - %0 = sparse_tensor.init [%c10, %c20] : tensor - return %0 : tensor -} - -// ----- - func.func @invalid_pointers_dense(%arg0: tensor<128xf64>) -> memref { %c = arith.constant 0 : index // expected-error@+1 {{expected a sparse tensor to get pointers}} diff --git a/mlir/test/Dialect/SparseTensor/roundtrip.mlir b/mlir/test/Dialect/SparseTensor/roundtrip.mlir --- a/mlir/test/Dialect/SparseTensor/roundtrip.mlir +++ b/mlir/test/Dialect/SparseTensor/roundtrip.mlir @@ -13,22 +13,6 @@ // ----- -#SparseMatrix = #sparse_tensor.encoding<{dimLevelType = ["compressed", "compressed"]}> - -// CHECK-LABEL: func @sparse_init() -// CHECK-DAG: %[[C16:.*]] = arith.constant 16 : index -// CHECK-DAG: %[[C32:.*]] = arith.constant 32 : index -// CHECK: %[[T:.*]] = sparse_tensor.init[%[[C16]], %[[C32]]] : tensor -// CHECK: return %[[T]] : tensor -func.func @sparse_init() -> tensor { - %d1 = arith.constant 16 : index - %d2 = arith.constant 32 : index - %0 = sparse_tensor.init [%d1, %d2] : tensor - return %0 : tensor -} - -// ----- - #SparseVector = #sparse_tensor.encoding<{dimLevelType = ["compressed"]}> // CHECK-LABEL: func @sparse_release( diff --git a/mlir/test/Dialect/SparseTensor/sparse_expand.mlir b/mlir/test/Dialect/SparseTensor/sparse_expand.mlir --- a/mlir/test/Dialect/SparseTensor/sparse_expand.mlir +++ b/mlir/test/Dialect/SparseTensor/sparse_expand.mlir @@ -49,7 +49,7 @@ func.func @kernel(%arga: tensor) -> tensor { %c0 = arith.constant 0 : index %n = tensor.dim %arga, %c0 : tensor - %v = sparse_tensor.init [%n] : tensor + %v = bufferization.alloc_tensor(%n) : tensor %0 = linalg.generic #rowsum ins(%arga: tensor) outs(%v: tensor) { diff --git a/mlir/test/Dialect/SparseTensor/sparse_fp_ops.mlir b/mlir/test/Dialect/SparseTensor/sparse_fp_ops.mlir --- a/mlir/test/Dialect/SparseTensor/sparse_fp_ops.mlir +++ b/mlir/test/Dialect/SparseTensor/sparse_fp_ops.mlir @@ -355,8 +355,7 @@ // CHECK-SAME: %[[VAL_0:.*]]: tensor<32xf64, #sparse_tensor.encoding<{{{.*}}}>> { // CHECK-DAG: %[[VAL_1:.*]] = arith.constant 0 : index // CHECK-DAG: %[[VAL_2:.*]] = arith.constant 1 : index -// CHECK-DAG: %[[VAL_3:.*]] = arith.constant 32 : index -// CHECK: %[[VAL_4:.*]] = sparse_tensor.init{{\[}}%[[VAL_3]]] : tensor<32xf64, #sparse_tensor.encoding<{{{.*}}}>> +// CHECK: %[[VAL_4:.*]] = bufferization.alloc_tensor() : tensor<32xf64, #sparse_tensor.encoding<{{{.*}}}>> // CHECK: %[[VAL_5:.*]] = sparse_tensor.pointers %[[VAL_0]], %[[VAL_1]] : tensor<32xf64, #sparse_tensor.encoding<{{{.*}}}>> to memref // CHECK: %[[VAL_6:.*]] = sparse_tensor.indices %[[VAL_0]], %[[VAL_1]] : tensor<32xf64, #sparse_tensor.encoding<{{{.*}}}>> to memref // CHECK: %[[VAL_7:.*]] = sparse_tensor.values %[[VAL_0]] : tensor<32xf64, #sparse_tensor.encoding<{{{.*}}}>> to memref @@ -382,7 +381,7 @@ // CHECK: } func.func @zero_preserving_math(%arga: tensor<32xf64, #SV>) -> tensor<32xf64, #SV> { %c32 = arith.constant 32 : index - %xinp = sparse_tensor.init [%c32] : tensor<32xf64, #SV> + %xinp = bufferization.alloc_tensor() : tensor<32xf64, #SV> %0 = linalg.generic #trait1 ins(%arga: tensor<32xf64, #SV>) outs(%xinp: tensor<32xf64, #SV>) { diff --git a/mlir/test/Dialect/SparseTensor/sparse_index.mlir b/mlir/test/Dialect/SparseTensor/sparse_index.mlir --- a/mlir/test/Dialect/SparseTensor/sparse_index.mlir +++ b/mlir/test/Dialect/SparseTensor/sparse_index.mlir @@ -24,7 +24,7 @@ // CHECK-DAG: %[[VAL_2:.*]] = arith.constant 1 : index // CHECK-DAG: %[[VAL_3:.*]] = tensor.dim %[[VAL_0]], %[[VAL_1]] : tensor %1 = tensor.dim %arga, %c1 : tensor - %init = sparse_tensor.init [%0, %1] : tensor + %init = bufferization.alloc_tensor(%0, %1) : tensor %r = linalg.generic #trait ins(%arga: tensor) outs(%init: tensor) { @@ -75,7 +75,7 @@ // CHECK-DAG: %[[VAL_3:.*]] = arith.constant 2 : index // CHECK-DAG: %[[VAL_4:.*]] = tensor.dim %[[VAL_0]], %[[VAL_1]] : tensor %1 = tensor.dim %arga, %c1 : tensor - %init = sparse_tensor.init [%0, %1] : tensor + %init = bufferization.alloc_tensor(%0, %1) : tensor %r = linalg.generic #trait ins(%arga: tensor) outs(%init: tensor) { diff --git a/mlir/test/Dialect/SparseTensor/sparse_kernels.mlir b/mlir/test/Dialect/SparseTensor/sparse_kernels.mlir --- a/mlir/test/Dialect/SparseTensor/sparse_kernels.mlir +++ b/mlir/test/Dialect/SparseTensor/sparse_kernels.mlir @@ -60,13 +60,12 @@ // CHECK-LABEL: func @matmul2( // CHECK-SAME: %[[VAL_0:.*]]: tensor<4x8xf64, #sparse_tensor.encoding<{{{.*}}}>>, // CHECK-SAME: %[[VAL_1:.*]]: tensor<8x4xf64, #sparse_tensor.encoding<{{{.*}}}>> { -// CHECK-DAG: %[[VAL_2:.*]] = arith.constant 4 : index // CHECK-DAG: %[[VAL_3:.*]] = arith.constant 0 : index // CHECK-DAG: %[[VAL_4:.*]] = arith.constant 1 : index // CHECK-DAG: %[[VAL_5:.*]] = arith.constant 2 : index // CHECK-DAG: %[[VAL_6:.*]] = arith.constant false // CHECK-DAG: %[[VAL_7:.*]] = arith.constant true -// CHECK: %[[VAL_8:.*]] = sparse_tensor.init{{\[}}%[[VAL_2]], %[[VAL_2]]] : tensor<4x4xf64, #sparse_tensor.encoding<{{{.*}}}>> +// CHECK: %[[VAL_8:.*]] = bufferization.alloc_tensor() : tensor<4x4xf64, #sparse_tensor.encoding<{{{.*}}}>> // CHECK: %[[VAL_9:.*]] = sparse_tensor.pointers %[[VAL_0]], %[[VAL_3]] : tensor<4x8xf64, #sparse_tensor.encoding<{{{.*}}}>> to memref // CHECK: %[[VAL_10:.*]] = sparse_tensor.indices %[[VAL_0]], %[[VAL_3]] : tensor<4x8xf64, #sparse_tensor.encoding<{{{.*}}}>> to memref // CHECK: %[[VAL_11:.*]] = sparse_tensor.pointers %[[VAL_0]], %[[VAL_4]] : tensor<4x8xf64, #sparse_tensor.encoding<{{{.*}}}>> to memref @@ -147,7 +146,7 @@ func.func @matmul2(%A: tensor<4x8xf64, #DCSR>, %B: tensor<8x4xf64, #DCSR>) -> tensor<4x4xf64, #DCSR> { %c4 = arith.constant 4 : index - %C = sparse_tensor.init [%c4, %c4] : tensor<4x4xf64, #DCSR> + %C = bufferization.alloc_tensor() : tensor<4x4xf64, #DCSR> %D = linalg.matmul ins(%A, %B: tensor<4x8xf64, #DCSR>, tensor<8x4xf64, #DCSR>) outs(%C: tensor<4x4xf64, #DCSR>) -> tensor<4x4xf64, #DCSR> diff --git a/mlir/test/Dialect/SparseTensor/sparse_out.mlir b/mlir/test/Dialect/SparseTensor/sparse_out.mlir --- a/mlir/test/Dialect/SparseTensor/sparse_out.mlir +++ b/mlir/test/Dialect/SparseTensor/sparse_out.mlir @@ -103,11 +103,10 @@ // CHECK-SAME: %[[VAL_0:.*]]: tensor<10x20xf32, #sparse_tensor.encoding<{{.*}}>> // CHECK-DAG: %[[VAL_1:.*]] = arith.constant 2.000000e+00 : f32 // CHECK-DAG: %[[VAL_2:.*]] = arith.constant 10 : index -// CHECK-DAG: %[[VAL_3:.*]] = arith.constant 20 : index // CHECK-DAG: %[[VAL_4:.*]] = arith.constant 1 : index // CHECK-DAG: %[[VAL_5:.*]] = arith.constant 2 : index // CHECK-DAG: %[[VAL_6:.*]] = arith.constant 0 : index -// CHECK: %[[VAL_7:.*]] = sparse_tensor.init{{\[}}%[[VAL_2]], %[[VAL_3]]] : tensor<10x20xf32, #sparse_tensor.encoding<{{.*}}>> +// CHECK: %[[VAL_7:.*]] = bufferization.alloc_tensor() : tensor<10x20xf32, #sparse_tensor.encoding<{{.*}}>> // CHECK: %[[VAL_8:.*]] = sparse_tensor.pointers %[[VAL_0]], %[[VAL_4]] : tensor<10x20xf32, #sparse_tensor.encoding<{{.*}}>> // CHECK: %[[VAL_9:.*]] = sparse_tensor.indices %[[VAL_0]], %[[VAL_4]] : tensor<10x20xf32, #sparse_tensor.encoding<{{.*}}>> // CHECK: %[[VAL_10:.*]] = sparse_tensor.values %[[VAL_0]] : tensor<10x20xf32, #sparse_tensor.encoding<{{.*}}>> @@ -130,9 +129,7 @@ // CHECK: } func.func @sparse_truly_dynamic(%arga: tensor<10x20xf32, #CSR>) -> tensor<10x20xf32, #DCSR> { %s = arith.constant 2.0 : f32 - %d10 = arith.constant 10 : index - %d20 = arith.constant 20 : index - %xm = sparse_tensor.init [%d10, %d20] : tensor<10x20xf32, #DCSR> + %xm = bufferization.alloc_tensor() : tensor<10x20xf32, #DCSR> %0 = linalg.generic #trait_scale ins(%arga: tensor<10x20xf32, #CSR>) outs(%xm: tensor<10x20xf32, #DCSR>) { @@ -162,7 +159,7 @@ // CHECK-DAG: %[[VAL_5:.*]] = arith.constant 0 : i32 // CHECK: %[[VAL_6:.*]] = tensor.dim %[[VAL_0]], %[[VAL_2]] : tensor> // CHECK: %[[VAL_7:.*]] = tensor.dim %[[VAL_0]], %[[VAL_3]] : tensor> -// CHECK: %[[VAL_8:.*]] = sparse_tensor.init{{\[}}%[[VAL_6]], %[[VAL_7]]] : tensor> +// CHECK: %[[VAL_8:.*]] = bufferization.alloc_tensor(%[[VAL_6]], %[[VAL_7]]) : tensor> // CHECK: %[[VAL_9:.*]] = sparse_tensor.pointers %[[VAL_0]], %[[VAL_2]] : tensor> to memref // CHECK: %[[VAL_10:.*]] = sparse_tensor.indices %[[VAL_0]], %[[VAL_2]] : tensor> to memref // CHECK: %[[VAL_11:.*]] = sparse_tensor.pointers %[[VAL_0]], %[[VAL_3]] : tensor> to memref @@ -288,7 +285,7 @@ %c1 = arith.constant 1 : index %d0 = tensor.dim %arga, %c0 : tensor %d1 = tensor.dim %arga, %c1 : tensor - %xinit = sparse_tensor.init [%d0, %d1] : tensor + %xinit = bufferization.alloc_tensor(%d0, %d1) : tensor %0 = linalg.generic #trait_sumred ins(%arga, %argb: tensor, tensor) @@ -321,7 +318,7 @@ // CHECK-DAG: %[[VAL_6:.*]] = arith.constant true // CHECK: %[[VAL_7:.*]] = tensor.dim %[[VAL_0]], %[[VAL_2]] : tensor> // CHECK: %[[VAL_8:.*]] = tensor.dim %[[VAL_1]], %[[VAL_3]] : tensor> -// CHECK: %[[VAL_9:.*]] = sparse_tensor.init{{\[}}%[[VAL_7]], %[[VAL_8]]] : tensor> +// CHECK: %[[VAL_9:.*]] = bufferization.alloc_tensor(%[[VAL_7]], %[[VAL_8]]) : tensor> // CHECK: %[[VAL_10:.*]] = sparse_tensor.pointers %[[VAL_0]], %[[VAL_2]] : tensor> to memref // CHECK: %[[VAL_11:.*]] = sparse_tensor.indices %[[VAL_0]], %[[VAL_2]] : tensor> to memref // CHECK: %[[VAL_12:.*]] = sparse_tensor.pointers %[[VAL_0]], %[[VAL_3]] : tensor> to memref @@ -405,7 +402,7 @@ %c1 = arith.constant 1 : index %d0 = tensor.dim %arga, %c0 : tensor %d1 = tensor.dim %argb, %c1 : tensor - %cinit = sparse_tensor.init [%d0, %d1] : tensor + %cinit = bufferization.alloc_tensor(%d0, %d1) : tensor %0 = linalg.generic #trait_matmat ins(%arga, %argb: tensor, tensor) diff --git a/mlir/test/Integration/Dialect/SparseTensor/CPU/dense_output.mlir b/mlir/test/Integration/Dialect/SparseTensor/CPU/dense_output.mlir --- a/mlir/test/Integration/Dialect/SparseTensor/CPU/dense_output.mlir +++ b/mlir/test/Integration/Dialect/SparseTensor/CPU/dense_output.mlir @@ -46,7 +46,7 @@ %c1 = arith.constant 1 : index %d0 = tensor.dim %arga, %c0 : tensor %d1 = tensor.dim %arga, %c1 : tensor - %init = sparse_tensor.init [%d0, %d1] : tensor + %init = bufferization.alloc_tensor(%d0, %d1) : tensor %0 = linalg.generic #trait_assign ins(%arga: tensor) outs(%init: tensor) { diff --git a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_binary.mlir b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_binary.mlir --- a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_binary.mlir +++ b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_binary.mlir @@ -42,7 +42,7 @@ %argb: tensor) -> tensor { %c = arith.constant 0 : index %d = tensor.dim %arga, %c : tensor - %xv = sparse_tensor.init [%d] : tensor + %xv = bufferization.alloc_tensor(%d) : tensor %0 = linalg.generic #trait_vec_op ins(%arga, %argb: tensor, tensor) outs(%xv: tensor) { @@ -67,7 +67,7 @@ %argb: tensor) -> tensor { %c = arith.constant 0 : index %d = tensor.dim %arga, %c : tensor - %xv = sparse_tensor.init [%d] : tensor + %xv = bufferization.alloc_tensor(%d) : tensor %0 = linalg.generic #trait_vec_op ins(%arga, %argb: tensor, tensor) outs(%xv: tensor) { @@ -91,7 +91,7 @@ %argb: tensor) -> tensor { %c = arith.constant 0 : index %d = tensor.dim %arga, %c : tensor - %xv = sparse_tensor.init [%d] : tensor + %xv = bufferization.alloc_tensor(%d) : tensor %0 = linalg.generic #trait_vec_op ins(%arga, %argb: tensor, tensor) outs(%xv: tensor) { @@ -109,7 +109,7 @@ func.func @vector_index(%arga: tensor) -> tensor { %c = arith.constant 0 : index %d = tensor.dim %arga, %c : tensor - %xv = sparse_tensor.init [%d] : tensor + %xv = bufferization.alloc_tensor(%d) : tensor %0 = linalg.generic #trait_vec_scale ins(%arga: tensor) outs(%xv: tensor) { @@ -136,7 +136,7 @@ %c1 = arith.constant 1 : index %d0 = tensor.dim %arga, %c0 : tensor %d1 = tensor.dim %arga, %c1 : tensor - %xv = sparse_tensor.init [%d0, %d1] : tensor + %xv = bufferization.alloc_tensor(%d0, %d1) : tensor %0 = linalg.generic #trait_mat_op ins(%arga, %argb: tensor, tensor) outs(%xv: tensor) { @@ -291,4 +291,4 @@ sparse_tensor.release %5 : tensor return } -} \ No newline at end of file +} diff --git a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_complex32.mlir b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_complex32.mlir --- a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_complex32.mlir +++ b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_complex32.mlir @@ -22,7 +22,7 @@ -> tensor, #SparseVector> { %c = arith.constant 0 : index %d = tensor.dim %arga, %c : tensor, #SparseVector> - %xv = sparse_tensor.init [%d] : tensor, #SparseVector> + %xv = bufferization.alloc_tensor(%d) : tensor, #SparseVector> %0 = linalg.generic #trait_op ins(%arga, %argb: tensor, #SparseVector>, tensor, #SparseVector>) @@ -39,7 +39,7 @@ -> tensor, #SparseVector> { %c = arith.constant 0 : index %d = tensor.dim %arga, %c : tensor, #SparseVector> - %xv = sparse_tensor.init [%d] : tensor, #SparseVector> + %xv = bufferization.alloc_tensor(%d) : tensor, #SparseVector> %0 = linalg.generic #trait_op ins(%arga, %argb: tensor, #SparseVector>, tensor, #SparseVector>) diff --git a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_complex64.mlir b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_complex64.mlir --- a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_complex64.mlir +++ b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_complex64.mlir @@ -22,7 +22,7 @@ -> tensor, #SparseVector> { %c = arith.constant 0 : index %d = tensor.dim %arga, %c : tensor, #SparseVector> - %xv = sparse_tensor.init [%d] : tensor, #SparseVector> + %xv = bufferization.alloc_tensor(%d) : tensor, #SparseVector> %0 = linalg.generic #trait_op ins(%arga, %argb: tensor, #SparseVector>, tensor, #SparseVector>) @@ -39,7 +39,7 @@ -> tensor, #SparseVector> { %c = arith.constant 0 : index %d = tensor.dim %arga, %c : tensor, #SparseVector> - %xv = sparse_tensor.init [%d] : tensor, #SparseVector> + %xv = bufferization.alloc_tensor(%d) : tensor, #SparseVector> %0 = linalg.generic #trait_op ins(%arga, %argb: tensor, #SparseVector>, tensor, #SparseVector>) diff --git a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_complex_ops.mlir b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_complex_ops.mlir --- a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_complex_ops.mlir +++ b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_complex_ops.mlir @@ -31,7 +31,7 @@ -> tensor, #SparseVector> { %c0 = arith.constant 0 : index %d = tensor.dim %arga, %c0 : tensor, #SparseVector> - %xv = sparse_tensor.init [%d] : tensor, #SparseVector> + %xv = bufferization.alloc_tensor(%d) : tensor, #SparseVector> %0 = linalg.generic #trait_op2 ins(%arga, %argb: tensor, #SparseVector>, tensor, #SparseVector>) @@ -48,7 +48,7 @@ -> tensor, #SparseVector> { %c0 = arith.constant 0 : index %d = tensor.dim %arga, %c0 : tensor, #SparseVector> - %xv = sparse_tensor.init [%d] : tensor, #SparseVector> + %xv = bufferization.alloc_tensor(%d) : tensor, #SparseVector> %0 = linalg.generic #trait_op1 ins(%arga: tensor, #SparseVector>) outs(%xv: tensor, #SparseVector>) { @@ -63,7 +63,7 @@ -> tensor, #SparseVector> { %c0 = arith.constant 0 : index %d = tensor.dim %arga, %c0 : tensor, #SparseVector> - %xv = sparse_tensor.init [%d] : tensor, #SparseVector> + %xv = bufferization.alloc_tensor(%d) : tensor, #SparseVector> %0 = linalg.generic #trait_op1 ins(%arga: tensor, #SparseVector>) outs(%xv: tensor, #SparseVector>) { @@ -78,7 +78,7 @@ -> tensor, #SparseVector> { %c0 = arith.constant 0 : index %d = tensor.dim %arga, %c0 : tensor, #SparseVector> - %xv = sparse_tensor.init [%d] : tensor, #SparseVector> + %xv = bufferization.alloc_tensor(%d) : tensor, #SparseVector> %0 = linalg.generic #trait_op1 ins(%arga: tensor, #SparseVector>) outs(%xv: tensor, #SparseVector>) { @@ -93,7 +93,7 @@ -> tensor, #SparseVector> { %c0 = arith.constant 0 : index %d = tensor.dim %arga, %c0 : tensor, #SparseVector> - %xv = sparse_tensor.init [%d] : tensor, #SparseVector> + %xv = bufferization.alloc_tensor(%d) : tensor, #SparseVector> %0 = linalg.generic #trait_op1 ins(%arga: tensor, #SparseVector>) outs(%xv: tensor, #SparseVector>) { @@ -111,7 +111,7 @@ -> tensor, #SparseVector> { %c0 = arith.constant 0 : index %d = tensor.dim %arga, %c0 : tensor, #SparseVector> - %xv = sparse_tensor.init [%d] : tensor, #SparseVector> + %xv = bufferization.alloc_tensor(%d) : tensor, #SparseVector> %c = complex.constant [2.0 : f64, 0.0 : f64] : complex %0 = linalg.generic #trait_op1 ins(%arga: tensor, #SparseVector>) @@ -127,7 +127,7 @@ -> tensor { %c0 = arith.constant 0 : index %d = tensor.dim %arga, %c0 : tensor, #SparseVector> - %xv = sparse_tensor.init [%d] : tensor + %xv = bufferization.alloc_tensor(%d) : tensor %0 = linalg.generic #trait_op1 ins(%arga: tensor, #SparseVector>) outs(%xv: tensor) { diff --git a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_index.mlir b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_index.mlir --- a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_index.mlir +++ b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_index.mlir @@ -40,8 +40,7 @@ // func.func @sparse_index_1d_conj(%arga: tensor<8xi64, #SparseVector>) -> tensor<8xi64, #SparseVector> { - %d0 = arith.constant 8 : index - %init = sparse_tensor.init [%d0] : tensor<8xi64, #SparseVector> + %init = bufferization.alloc_tensor() : tensor<8xi64, #SparseVector> %r = linalg.generic #trait_1d ins(%arga: tensor<8xi64, #SparseVector>) outs(%init: tensor<8xi64, #SparseVector>) { @@ -59,8 +58,7 @@ // func.func @sparse_index_1d_disj(%arga: tensor<8xi64, #SparseVector>) -> tensor<8xi64, #SparseVector> { - %d0 = arith.constant 8 : index - %init = sparse_tensor.init [%d0] : tensor<8xi64, #SparseVector> + %init = bufferization.alloc_tensor() : tensor<8xi64, #SparseVector> %r = linalg.generic #trait_1d ins(%arga: tensor<8xi64, #SparseVector>) outs(%init: tensor<8xi64, #SparseVector>) { @@ -78,9 +76,7 @@ // func.func @sparse_index_2d_conj(%arga: tensor<3x4xi64, #SparseMatrix>) -> tensor<3x4xi64, #SparseMatrix> { - %d0 = arith.constant 3 : index - %d1 = arith.constant 4 : index - %init = sparse_tensor.init [%d0, %d1] : tensor<3x4xi64, #SparseMatrix> + %init = bufferization.alloc_tensor() : tensor<3x4xi64, #SparseMatrix> %r = linalg.generic #trait_2d ins(%arga: tensor<3x4xi64, #SparseMatrix>) outs(%init: tensor<3x4xi64, #SparseMatrix>) { @@ -101,9 +97,7 @@ // func.func @sparse_index_2d_disj(%arga: tensor<3x4xi64, #SparseMatrix>) -> tensor<3x4xi64, #SparseMatrix> { - %d0 = arith.constant 3 : index - %d1 = arith.constant 4 : index - %init = sparse_tensor.init [%d0, %d1] : tensor<3x4xi64, #SparseMatrix> + %init = bufferization.alloc_tensor() : tensor<3x4xi64, #SparseMatrix> %r = linalg.generic #trait_2d ins(%arga: tensor<3x4xi64, #SparseMatrix>) outs(%init: tensor<3x4xi64, #SparseMatrix>) { @@ -121,9 +115,7 @@ func.func @add_outer_2d(%arg0: tensor<2x3xf32, #SparseMatrix>) -> tensor<2x3xf32, #SparseMatrix> { - %c2 = arith.constant 2 : index - %c3 = arith.constant 3 : index - %0 = sparse_tensor.init[%c2, %c3] : tensor<2x3xf32, #SparseMatrix> + %0 = bufferization.alloc_tensor() : tensor<2x3xf32, #SparseMatrix> %1 = linalg.generic #trait_2d ins(%arg0 : tensor<2x3xf32, #SparseMatrix>) outs(%0 : tensor<2x3xf32, #SparseMatrix>) { diff --git a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_matmul.mlir b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_matmul.mlir --- a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_matmul.mlir +++ b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_matmul.mlir @@ -31,8 +31,7 @@ // func.func @matmul2(%A: tensor<4x8xf64, #CSR>, %B: tensor<8x4xf64, #CSR>) -> tensor<4x4xf64, #CSR> { - %c4 = arith.constant 4 : index - %C = sparse_tensor.init [%c4, %c4] : tensor<4x4xf64, #CSR> + %C = bufferization.alloc_tensor() : tensor<4x4xf64, #CSR> %D = linalg.matmul ins(%A, %B: tensor<4x8xf64, #CSR>, tensor<8x4xf64, #CSR>) outs(%C: tensor<4x4xf64, #CSR>) -> tensor<4x4xf64, #CSR> @@ -44,8 +43,7 @@ // func.func @matmul3(%A: tensor<4x8xf64, #DCSR>, %B: tensor<8x4xf64, #DCSR>) -> tensor<4x4xf64, #DCSR> { - %c4 = arith.constant 4 : index - %C = sparse_tensor.init [%c4, %c4] : tensor<4x4xf64, #DCSR> + %C = bufferization.alloc_tensor() : tensor<4x4xf64, #DCSR> %D = linalg.matmul ins(%A, %B: tensor<4x8xf64, #DCSR>, tensor<8x4xf64, #DCSR>) outs(%C: tensor<4x4xf64, #DCSR>) -> tensor<4x4xf64, #DCSR> diff --git a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_matrix_ops.mlir b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_matrix_ops.mlir --- a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_matrix_ops.mlir +++ b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_matrix_ops.mlir @@ -42,7 +42,7 @@ %c1 = arith.constant 1 : index %d0 = tensor.dim %arga, %c0 : tensor %d1 = tensor.dim %arga, %c1 : tensor - %xm = sparse_tensor.init [%d0, %d1] : tensor + %xm = bufferization.alloc_tensor(%d0, %d1) : tensor %0 = linalg.generic #trait_scale ins(%arga: tensor) outs(%xm: tensor) { @@ -73,7 +73,7 @@ %c1 = arith.constant 1 : index %d0 = tensor.dim %arga, %c0 : tensor %d1 = tensor.dim %arga, %c1 : tensor - %xv = sparse_tensor.init [%d0, %d1] : tensor + %xv = bufferization.alloc_tensor(%d0, %d1) : tensor %0 = linalg.generic #trait_op ins(%arga, %argb: tensor, tensor) outs(%xv: tensor) { @@ -91,7 +91,7 @@ %c1 = arith.constant 1 : index %d0 = tensor.dim %arga, %c0 : tensor %d1 = tensor.dim %arga, %c1 : tensor - %xv = sparse_tensor.init [%d0, %d1] : tensor + %xv = bufferization.alloc_tensor(%d0, %d1) : tensor %0 = linalg.generic #trait_op ins(%arga, %argb: tensor, tensor) outs(%xv: tensor) { diff --git a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_out_mult_elt.mlir b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_out_mult_elt.mlir --- a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_out_mult_elt.mlir +++ b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_out_mult_elt.mlir @@ -22,9 +22,7 @@ // Sparse kernel. func.func @sparse_mult_elt( %arga: tensor<32x16xf32, #DCSR>, %argb: tensor<32x16xf32, #DCSR>) -> tensor<32x16xf32, #DCSR> { - %c16 = arith.constant 16 : index - %c32 = arith.constant 32 : index - %argx = sparse_tensor.init [%c32, %c16] : tensor<32x16xf32, #DCSR> + %argx = bufferization.alloc_tensor() : tensor<32x16xf32, #DCSR> %0 = linalg.generic #trait_mult_elt ins(%arga, %argb: tensor<32x16xf32, #DCSR>, tensor<32x16xf32, #DCSR>) outs(%argx: tensor<32x16xf32, #DCSR>) { diff --git a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_out_reduction.mlir b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_out_reduction.mlir --- a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_out_reduction.mlir +++ b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_out_reduction.mlir @@ -30,7 +30,7 @@ %c1 = arith.constant 1 : index %d0 = tensor.dim %arga, %c0 : tensor %d1 = tensor.dim %arga, %c1 : tensor - %xinit = sparse_tensor.init [%d0, %d1] : tensor + %xinit = bufferization.alloc_tensor(%d0, %d1): tensor %0 = linalg.generic #redsum ins(%arga, %argb: tensor, tensor) diff --git a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_re_im.mlir b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_re_im.mlir --- a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_re_im.mlir +++ b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_re_im.mlir @@ -20,7 +20,7 @@ -> tensor { %c = arith.constant 0 : index %d = tensor.dim %arga, %c : tensor, #SparseVector> - %xv = sparse_tensor.init [%d] : tensor + %xv = bufferization.alloc_tensor(%d) : tensor %0 = linalg.generic #trait_op ins(%arga: tensor, #SparseVector>) outs(%xv: tensor) { @@ -35,7 +35,7 @@ -> tensor { %c = arith.constant 0 : index %d = tensor.dim %arga, %c : tensor, #SparseVector> - %xv = sparse_tensor.init [%d] : tensor + %xv = bufferization.alloc_tensor(%d) : tensor %0 = linalg.generic #trait_op ins(%arga: tensor, #SparseVector>) outs(%xv: tensor) { diff --git a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_sampled_mm_fusion.mlir b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_sampled_mm_fusion.mlir --- a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_sampled_mm_fusion.mlir +++ b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_sampled_mm_fusion.mlir @@ -101,8 +101,7 @@ func.func @sparse_sampled_dd(%args: tensor<8x8xf64, #SM>, %arga: tensor<8x8xf64>, %argb: tensor<8x8xf64>) -> tensor<8x8xf64, #SM> { - %c8 = arith.constant 8 : index - %1 = sparse_tensor.init [%c8, %c8] : tensor<8x8xf64, #SM> + %1 = bufferization.alloc_tensor() : tensor<8x8xf64, #SM> %2 = linalg.generic #trait_sampled_dense_dense ins(%args, %arga, %argb: tensor<8x8xf64, #SM>, tensor<8x8xf64>, tensor<8x8xf64>) @@ -135,8 +134,7 @@ linalg.yield %q : f64 } -> tensor<8x8xf64> // Sample the result with elements-wise multiplication with sparse matrix. - %c8 = arith.constant 8 : index - %3 = sparse_tensor.init [%c8, %c8] : tensor<8x8xf64, #SM> + %3 = bufferization.alloc_tensor() : tensor<8x8xf64, #SM> %4 = linalg.generic #trait_scale ins(%2, %args : tensor<8x8xf64>, tensor<8x8xf64, #SM>) outs(%3 : tensor<8x8xf64, #SM>) { diff --git a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_sign.mlir b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_sign.mlir --- a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_sign.mlir +++ b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_sign.mlir @@ -30,7 +30,7 @@ -> tensor { %c0 = arith.constant 0 : index %d = tensor.dim %arg0, %c0 : tensor - %xin = sparse_tensor.init [%d] : tensor + %xin = bufferization.alloc_tensor(%d) : tensor %0 = linalg.generic #trait_op ins(%arg0: tensor) outs(%xin: tensor) { diff --git a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_tensor_ops.mlir b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_tensor_ops.mlir --- a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_tensor_ops.mlir +++ b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_tensor_ops.mlir @@ -29,7 +29,7 @@ %d0 = tensor.dim %arga, %c0 : tensor %d1 = tensor.dim %arga, %c1 : tensor %d2 = tensor.dim %arga, %c2 : tensor - %xm = sparse_tensor.init [%d0, %d1, %d2] : tensor + %xm = bufferization.alloc_tensor(%d0, %d1, %d2) : tensor %0 = linalg.generic #trait_scale ins(%arga: tensor) outs(%xm: tensor) { diff --git a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_transpose.mlir b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_transpose.mlir --- a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_transpose.mlir +++ b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_transpose.mlir @@ -32,9 +32,7 @@ func.func @sparse_transpose(%arga: tensor<3x4xf64, #DCSR>) -> tensor<4x3xf64, #DCSR> { %t = sparse_tensor.convert %arga : tensor<3x4xf64, #DCSR> to tensor<3x4xf64, #DCSC> - %c3 = arith.constant 3 : index - %c4 = arith.constant 4 : index - %i = sparse_tensor.init [%c4, %c3] : tensor<4x3xf64, #DCSR> + %i = bufferization.alloc_tensor() : tensor<4x3xf64, #DCSR> %0 = linalg.generic #transpose_trait ins(%t: tensor<3x4xf64, #DCSC>) diff --git a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_unary.mlir b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_unary.mlir --- a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_unary.mlir +++ b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_unary.mlir @@ -32,7 +32,7 @@ %c = arith.constant 0 : index %ci1 = arith.constant 1 : i32 %d = tensor.dim %arga, %c : tensor - %xv = sparse_tensor.init [%d] : tensor + %xv = bufferization.alloc_tensor(%d) : tensor %0 = linalg.generic #trait_vec_scale ins(%arga: tensor) outs(%xv: tensor) { @@ -52,7 +52,7 @@ %c = arith.constant 0 : index %cf1 = arith.constant 1.0 : f64 %d = tensor.dim %arga, %c : tensor - %xv = sparse_tensor.init [%d] : tensor + %xv = bufferization.alloc_tensor(%d) : tensor %0 = linalg.generic #trait_vec_scale ins(%arga: tensor) outs(%xv: tensor) { @@ -79,7 +79,7 @@ %cfmax = arith.constant 7.0 : f64 %d0 = tensor.dim %argx, %c0 : tensor %d1 = tensor.dim %argx, %c1 : tensor - %xv = sparse_tensor.init [%d0, %d1] : tensor + %xv = bufferization.alloc_tensor(%d0, %d1) : tensor %0 = linalg.generic #trait_mat_scale ins(%argx: tensor) outs(%xv: tensor) { @@ -202,4 +202,4 @@ sparse_tensor.release %2 : tensor return } -} \ No newline at end of file +} diff --git a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_vector_ops.mlir b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_vector_ops.mlir --- a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_vector_ops.mlir +++ b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_vector_ops.mlir @@ -50,7 +50,7 @@ %s = arith.constant 2.0 : f64 %c = arith.constant 0 : index %d = tensor.dim %arga, %c : tensor - %xv = sparse_tensor.init [%d] : tensor + %xv = bufferization.alloc_tensor(%d) : tensor %0 = linalg.generic #trait_scale ins(%arga: tensor) outs(%xv: tensor) { @@ -79,7 +79,7 @@ %argb: tensor) -> tensor { %c = arith.constant 0 : index %d = tensor.dim %arga, %c : tensor - %xv = sparse_tensor.init [%d] : tensor + %xv = bufferization.alloc_tensor(%d) : tensor %0 = linalg.generic #trait_op ins(%arga, %argb: tensor, tensor) outs(%xv: tensor) { @@ -95,7 +95,7 @@ %argb: tensor) -> tensor { %c = arith.constant 0 : index %d = tensor.dim %arga, %c : tensor - %xv = sparse_tensor.init [%d] : tensor + %xv = bufferization.alloc_tensor(%d) : tensor %0 = linalg.generic #trait_op ins(%arga, %argb: tensor, tensor) outs(%xv: tensor) { @@ -111,7 +111,7 @@ %argb: tensor) -> tensor { %c = arith.constant 0 : index %d = tensor.dim %arga, %c : tensor - %xv = sparse_tensor.init [%d] : tensor + %xv = bufferization.alloc_tensor(%d) : tensor %0 = linalg.generic #trait_op ins(%arga, %argb: tensor, tensor) outs(%xv: tensor) { diff --git a/mlir/test/Integration/Dialect/SparseTensor/python/test_elementwise_add_sparse_output.py b/mlir/test/Integration/Dialect/SparseTensor/python/test_elementwise_add_sparse_output.py --- a/mlir/test/Integration/Dialect/SparseTensor/python/test_elementwise_add_sparse_output.py +++ b/mlir/test/Integration/Dialect/SparseTensor/python/test_elementwise_add_sparse_output.py @@ -35,9 +35,7 @@ func.func @sparse_add_elt( %arga: tensor<3x4xf64, #DCSR>, %argb: tensor<3x4xf64, #DCSR>) -> tensor<3x4xf64, #DCSR> { - %c3 = arith.constant 3 : index - %c4 = arith.constant 4 : index - %argx = sparse_tensor.init [%c3, %c4] : tensor<3x4xf64, #DCSR> + %argx = bufferization.alloc_tensor() : tensor<3x4xf64, #DCSR> %0 = linalg.generic #trait_add_elt ins(%arga, %argb: tensor<3x4xf64, #DCSR>, tensor<3x4xf64, #DCSR>) outs(%argx: tensor<3x4xf64, #DCSR>) { diff --git a/mlir/test/Integration/Dialect/SparseTensor/taco/tools/mlir_pytaco.py b/mlir/test/Integration/Dialect/SparseTensor/taco/tools/mlir_pytaco.py --- a/mlir/test/Integration/Dialect/SparseTensor/taco/tools/mlir_pytaco.py +++ b/mlir/test/Integration/Dialect/SparseTensor/taco/tools/mlir_pytaco.py @@ -34,6 +34,7 @@ from mlir import ir from mlir import runtime from mlir.dialects import arith +from mlir.dialects import bufferization from mlir.dialects import builtin from mlir.dialects import func from mlir.dialects import linalg @@ -889,8 +890,7 @@ mlir_type = _mlir_tensor_type(self.dst_dtype, self.dst_dims, self.dst_format.mlir_tensor_attr()) index_type = ir.IndexType.get() - dims = [arith.ConstantOp(index_type, d).result for d in mlir_type.shape] - return sparse_tensor.InitOp(mlir_type, dims) + return bufferization.AllocTensorOp(mlir_type, []) class _Stats: