diff --git a/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorCodegen.cpp b/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorCodegen.cpp --- a/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorCodegen.cpp +++ b/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorCodegen.cpp @@ -1234,6 +1234,134 @@ } }; +struct SparseNewOpConverter : public OpConversionPattern { + using OpConversionPattern::OpConversionPattern; + LogicalResult + matchAndRewrite(NewOp op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override { + Location loc = op.getLoc(); + const auto dstTp = getSparseTensorType(op.getResult()); + const auto encDst = dstTp.getEncoding(); + if (!dstTp.hasEncoding() || getCOOStart(encDst) != 0) + return failure(); + + // Implement the NewCOO(filename) as follows: + // reader = getSparseTensorReader(filename) + // nse = getSparseTensorNNZ() + // tmp = bufferization.alloc_tensor an ordered COO with + // dst dim ordering, size_hint = nse + // indices = to_indices_buffer(tmp) + // values = to_values(tmp) + // isSorted = getSparseTensorReaderRead(indices, values, dimOrdering) + // if!isSorted sort_coo(nse, indices, values) + // update storage specifier + // dst = sparse_tensor.ConvertOp tmp + + // Create a sparse tensor reader. + Value fileName = op.getSource(); + Type opaqueTp = getOpaquePointerType(rewriter); + Value reader = createFuncCall(rewriter, loc, "createSparseTensorReader", + {opaqueTp}, {fileName}, EmitCInterface::Off) + .getResult(0); + + Type indexTp = rewriter.getIndexType(); + const Dimension dimRank = dstTp.getDimRank(); + + // If the result tensor has dynamic dimensions, get the dynamic sizes from + // the sparse tensor reader. + SmallVector dynSizes; + if (dstTp.hasDynamicDimShape()) { + Value dimSizes = genAlloca(rewriter, loc, dimRank, indexTp); + createFuncCall(rewriter, loc, "copySparseTensorReaderDimSizes", {}, + {reader, dimSizes}, EmitCInterface::On) + .getResult(0); + ArrayRef dstShape = dstTp.getRankedTensorType().getShape(); + for (auto &d : llvm::enumerate(dstShape)) { + if (d.value() == ShapedType::kDynamic) { + dynSizes.push_back(rewriter.create( + loc, dimSizes, constantIndex(rewriter, loc, d.index()))); + } + } + } + + Value nse = createFuncCall(rewriter, loc, "getSparseTensorReaderNNZ", + {indexTp}, {reader}, EmitCInterface::Off) + .getResult(0); + // Construct allocation for each field. + SmallVector fields; + createAllocFields(rewriter, loc, dstTp, dynSizes, /*enableInit=*/false, + fields, nse); + MutSparseTensorDescriptor desc(dstTp, fields); + + // Read the COO tensor data. + Type eltTp = dstTp.getElementType(); + Type indBufEleTp = getIndexOverheadType(rewriter, encDst); + SmallString<32> getNextFuncName{"getSparseTensorReaderRead", + overheadTypeFunctionSuffix(indBufEleTp), + primaryTypeFunctionSuffix(eltTp)}; + + Value xs = desc.getAOSMemRef(); + Value ys = desc.getValMemRef(); + SmallVector dim2lvlValues(dimRank, Value()); + if (auto dimOrder = encDst.getDimOrdering()) { + assert(dimOrder.isPermutation() && "Got non-permutation"); + for (uint64_t l = 0; l < dimRank; l++) { + uint64_t d = dimOrder.getDimPosition(l); + dim2lvlValues[d] = constantIndex(rewriter, loc, l); + } + } else { + for (uint64_t l = 0; l < dimRank; l++) + dim2lvlValues[l] = constantIndex(rewriter, loc, l); + } + Value dim2lvl = allocaBuffer(rewriter, loc, dim2lvlValues); + + Value f = constantI1(rewriter, loc, false); + Value isSorted = + createFuncCall(rewriter, loc, getNextFuncName, {f.getType()}, + {reader, dim2lvl, xs, ys}, EmitCInterface::On) + .getResult(0); + + if (encDst.isOrderedLvl(dimRank - 1)) { + // Sort the COO tensor data, if it is not sorted yet. + Value notSorted = rewriter.create( + loc, arith::CmpIPredicate::eq, isSorted, f); + scf::IfOp ifOp = + rewriter.create(loc, notSorted, /*else*/ false); + rewriter.setInsertionPointToStart(&ifOp.getThenRegion().front()); + rewriter.create( + loc, nse, xs, ValueRange{ys}, rewriter.getIndexAttr(dimRank), + rewriter.getIndexAttr(0), SparseTensorSortKind::HybridQuickSort); + rewriter.setInsertionPointAfter(ifOp); + } + + // Set PtrMemRef0[1] = nse. + Value c1 = constantIndex(rewriter, loc, 1); + Value ptrMemref0 = desc.getPtrMemRef(0); + Type ptrEleTy = getMemRefType(ptrMemref0).getElementType(); + Value ptrNse = + ptrEleTy == nse.getType() + ? nse + : rewriter.create(loc, ptrEleTy, nse); + rewriter.create(loc, ptrNse, ptrMemref0, c1); + + // Update storage specifier. + Value idxSize = rewriter.create( + loc, nse, constantIndex(rewriter, loc, dimRank)); + desc.setSpecifierField(rewriter, loc, StorageSpecifierKind::IdxMemSize, 0, + idxSize); + desc.setSpecifierField(rewriter, loc, StorageSpecifierKind::ValMemSize, + std::nullopt, nse); + + // Release the sparse tensor reader. + createFuncCall(rewriter, loc, "delSparseTensorReader", {}, {reader}, + EmitCInterface::Off); + + // Replace operation with resulting memrefs. + rewriter.replaceOp(op, genTuple(rewriter, loc, dstTp, fields)); + return success(); + } +}; + } // namespace //===----------------------------------------------------------------------===// @@ -1253,8 +1381,8 @@ SparseInsertConverter, SparseToPointersConverter, SparseToIndicesConverter, SparseToIndicesBufferConverter, SparseToValuesConverter, SparseConvertConverter, - SparseNumberOfEntriesConverter>(typeConverter, - patterns.getContext()); + SparseNewOpConverter, SparseNumberOfEntriesConverter>( + typeConverter, patterns.getContext()); patterns.add(typeConverter, patterns.getContext(), enableBufferInitialization); } diff --git a/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorRewriting.cpp b/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorRewriting.cpp --- a/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorRewriting.cpp +++ b/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorRewriting.cpp @@ -999,92 +999,21 @@ Location loc = op.getLoc(); const auto dstTp = getSparseTensorType(op.getResult()); const auto encDst = dstTp.getEncoding(); - if (!dstTp.hasEncoding()) + if (!dstTp.hasEncoding() || getCOOStart(encDst) == 0) return failure(); - // Create a sparse tensor reader. - Value fileName = op.getSource(); - Type opaqueTp = getOpaquePointerType(rewriter); - Value reader = createFuncCall(rewriter, loc, "createSparseTensorReader", - {opaqueTp}, {fileName}, EmitCInterface::Off) - .getResult(0); - - // Allocate a temporary buffer for storing dimension sizes and indices. - Type indexTp = rewriter.getIndexType(); - const Dimension dimRank = dstTp.getDimRank(); - Value dimSizes = genAlloca(rewriter, loc, dimRank, indexTp); - - // If the result tensor has dynamic dimensions, get the dynamic sizes from - // the sparse tensor reader. - SmallVector dynSizesArray; - if (dstTp.hasDynamicDimShape()) { - createFuncCall(rewriter, loc, "copySparseTensorReaderDimSizes", {}, - {reader, dimSizes}, EmitCInterface::On) - .getResult(0); - ArrayRef dstShape = dstTp.getRankedTensorType().getShape(); - for (auto &d : llvm::enumerate(dstShape)) { - if (d.value() == ShapedType::kDynamic) { - dynSizesArray.push_back(rewriter.create( - loc, dimSizes, constantIndex(rewriter, loc, d.index()))); - } - } - } - // Implement the NewOp as follows: - // %tmp = bufferization.alloc_tensor : an unordered COO with identity - // storage ordering - // for i = 0 to nnz - // get the next element from the input file - // insert the element to %tmp - // %t = sparse_tensor.ConvertOp %tmp - Value c0 = constantIndex(rewriter, loc, 0); - Value c1 = constantIndex(rewriter, loc, 1); - Value nnz = createFuncCall(rewriter, loc, "getSparseTensorReaderNNZ", - {indexTp}, {reader}, EmitCInterface::Off) - .getResult(0); + // %orderedCoo = sparse_tensor.new %filename + // %t = sparse_tensor.ConvertOp %orderedCoo RankedTensorType cooTp = - getUnorderedCOOFromTypeWithOrdering(dstTp, dstTp.getDimToLvlMap()); - Value cooBuffer = - rewriter - .create(loc, cooTp, dynSizesArray, Value(), - /*sizeHint=*/nnz, Attribute()) - .getResult(); - - Type eltTp = dstTp.getElementType(); - Value value = genAllocaScalar(rewriter, loc, eltTp); - scf::ForOp forOp = rewriter.create(loc, c0, nnz, c1, - ArrayRef(cooBuffer)); - rewriter.setInsertionPointToStart(forOp.getBody()); - - SmallString<29> getNextFuncName{"getSparseTensorReaderNext", - primaryTypeFunctionSuffix(eltTp)}; - Value indices = dimSizes; // Reuse the indices memref to store indices. - createFuncCall(rewriter, loc, getNextFuncName, {}, {reader, indices, value}, - EmitCInterface::On); - SmallVector indicesArray(dimRank); - for (Dimension d = 0; d < dimRank; d++) { - // FIXME: `toStoredDim` is deprecated - indicesArray[toStoredDim(encDst, d)] = rewriter.create( - loc, indices, constantIndex(rewriter, loc, d)); - } - Value v = rewriter.create(loc, value); - Value t = rewriter.create(loc, v, forOp.getRegionIterArg(0), - indicesArray); - rewriter.create(loc, ArrayRef(t)); - rewriter.setInsertionPointAfter(forOp); - // Link SSA chain. - cooBuffer = forOp.getResult(0); - - // Release the sparse tensor reader. - createFuncCall(rewriter, loc, "delSparseTensorReader", {}, {reader}, - EmitCInterface::Off); - cooBuffer = rewriter.create(loc, cooBuffer, true); - Value newOp = rewriter.replaceOpWithNewOp( - op, dstTp.getRankedTensorType(), cooBuffer); - - // Release the unordered COO tensor buffer. - rewriter.setInsertionPointAfterValue(newOp); - rewriter.create(loc, cooBuffer); + getCOOFromTypeWithOrdering(dstTp, encDst.getDimOrdering(), true); + Value cooTensor = rewriter.create(loc, cooTp, op.getSource()); + Value convert = rewriter.replaceOpWithNewOp( + op, dstTp.getRankedTensorType(), cooTensor); + + // Release the ordered COO tensor. + rewriter.setInsertionPointAfterValue(convert); + rewriter.create(loc, cooTensor); return success(); } diff --git a/mlir/test/Dialect/SparseTensor/codegen.mlir b/mlir/test/Dialect/SparseTensor/codegen.mlir --- a/mlir/test/Dialect/SparseTensor/codegen.mlir +++ b/mlir/test/Dialect/SparseTensor/codegen.mlir @@ -50,6 +50,11 @@ dimLevelType = [ "compressed-nu", "singleton" ] }> +#CooPNo = #sparse_tensor.encoding<{ + dimLevelType = [ "compressed-nu", "singleton-no" ], + dimOrdering = affine_map<(i, j) -> (j, i)> +}> + #ccoo = #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "compressed-nu", "singleton" ] }> @@ -668,3 +673,102 @@ %0 = sparse_tensor.convert %arg0 : tensor<32xf32, #SparseVector> to tensor return %0 : tensor } + +// CHECK-LABEL: func.func @sparse_new_coo( +// CHECK-SAME: %[[A0:.*]]: !llvm.ptr) -> (memref, memref, memref, !sparse_tensor.storage_specifier<#sparse_tensor.encoding<{ dimLevelType = [ "compressed", "singleton" ] }>>) { +// CHECK-DAG: %[[A1:.*]] = arith.constant false +// CHECK-DAG: %[[A2:.*]] = arith.constant 1 : index +// CHECK-DAG: %[[A3:.*]] = arith.constant 0 : index +// CHECK-DAG: %[[A4:.*]] = arith.constant 2 : index +// CHECK: %[[A5:.*]] = call @createSparseTensorReader(%[[A0]]) +// CHECK: %[[A6:.*]] = memref.alloca() : memref<2xindex> +// CHECK: %[[A7:.*]] = memref.cast %[[A6]] : memref<2xindex> to memref +// CHECK: call @copySparseTensorReaderDimSizes(%[[A5]], %[[A7]]) : (!llvm.ptr, memref) -> () +// CHECK: %[[A8:.*]] = memref.load %[[A6]]{{\[}}%[[A3]]] : memref<2xindex> +// CHECK: %[[A9:.*]] = memref.load %[[A6]]{{\[}}%[[A2]]] : memref<2xindex> +// CHECK: %[[A10:.*]] = call @getSparseTensorReaderNNZ(%[[A5]]) +// CHECK: %[[A11:.*]] = arith.muli %[[A10]], %[[A4]] : index +// CHECK: %[[A12:.*]] = memref.alloc() : memref<2xindex> +// CHECK: %[[A13:.*]] = memref.cast %[[A12]] : memref<2xindex> to memref +// CHECK: %[[A14:.*]] = memref.alloc(%[[A11]]) : memref +// CHECK: %[[A15:.*]] = memref.alloc(%[[A10]]) : memref +// CHECK: %[[A16:.*]] = sparse_tensor.storage_specifier.init : !sparse_tensor.storage_specifier<#sparse_tensor.encoding<{ dimLevelType = [ "compressed", "singleton" ] }>> +// CHECK: %[[A17:.*]] = arith.index_cast %[[A8]] : index to i64 +// CHECK: %[[A18:.*]] = sparse_tensor.storage_specifier.set %[[A16]] dim_sz at 0 with %[[A17]] +// CHECK: %[[A19:.*]] = sparse_tensor.storage_specifier.get %[[A18]] ptr_mem_sz at 0 +// CHECK: %[[A20:.*]] = arith.index_cast %[[A19]] : i64 to index +// CHECK: %[[A21:.*]], %[[A22:.*]] = sparse_tensor.push_back %[[A20]], %[[A13]], %[[A3]] +// CHECK: %[[A23:.*]] = arith.index_cast %[[A22]] : index to i64 +// CHECK: %[[A24:.*]] = sparse_tensor.storage_specifier.set %[[A18]] ptr_mem_sz at 0 with %[[A23]] +// CHECK: %[[A25:.*]] = arith.index_cast %[[A9]] : index to i64 +// CHECK: %[[A26:.*]] = sparse_tensor.storage_specifier.set %[[A24]] dim_sz at 1 with %[[A25]] +// CHECK: %[[A27:.*]], %[[A28:.*]] = sparse_tensor.push_back %[[A22]], %[[A21]], %[[A3]], %[[A2]] +// CHECK: %[[A29:.*]] = arith.index_cast %[[A28]] : index to i64 +// CHECK: %[[A30:.*]] = sparse_tensor.storage_specifier.set %[[A26]] ptr_mem_sz at 0 with %[[A29]] +// CHECK: %[[A31:.*]] = memref.alloca() : memref<2xindex> +// CHECK: %[[A32:.*]] = memref.cast %[[A31]] : memref<2xindex> to memref +// CHECK: memref.store %[[A3]], %[[A31]]{{\[}}%[[A3]]] : memref<2xindex> +// CHECK: memref.store %[[A2]], %[[A31]]{{\[}}%[[A2]]] : memref<2xindex> +// CHECK: %[[A33:.*]] = call @getSparseTensorReaderRead0F32(%[[A5]], %[[A32]], %[[A14]], %[[A15]]) +// CHECK: %[[A34:.*]] = arith.cmpi eq, %[[A33]], %[[A1]] : i1 +// CHECK: scf.if %[[A34]] { +// CHECK: sparse_tensor.sort_coo hybrid_quick_sort %[[A10]], %[[A14]] jointly %[[A15]] {nx = 2 : index, ny = 0 : index} : memref jointly memref +// CHECK: } +// CHECK: memref.store %[[A10]], %[[A27]]{{\[}}%[[A2]]] : memref +// CHECK: %[[A35:.*]] = arith.index_cast %[[A11]] : index to i64 +// CHECK: %[[A36:.*]] = sparse_tensor.storage_specifier.set %[[A30]] idx_mem_sz at 0 with %[[A35]] +// CHECK: %[[A37:.*]] = arith.index_cast %[[A10]] : index to i64 +// CHECK: %[[A38:.*]] = sparse_tensor.storage_specifier.set %[[A36]] val_mem_sz with %[[A37]] +// CHECK: call @delSparseTensorReader(%[[A5]]) : (!llvm.ptr) -> () +// CHECK: return %[[A27]], %[[A14]], %[[A15]], %[[A38]] +func.func @sparse_new_coo(%arg0: !llvm.ptr) -> tensor { + %0 = sparse_tensor.new %arg0 : !llvm.ptr to tensor + return %0 : tensor +} + +// CHECK-LABEL: func.func @sparse_new_coo_permute_no( +// CHECK-SAME: %[[A0:.*]]: !llvm.ptr) -> (memref, memref, memref, !sparse_tensor.storage_specifier<#sparse_tensor.encoding<{ dimLevelType = [ "compressed", "singleton" ] }>>) { +// CHECK-DAG: %[[A1:.*]] = arith.constant 1 : index +// CHECK-DAG: %[[A2:.*]] = arith.constant 0 : index +// CHECK-DAG: %[[A3:.*]] = arith.constant 2 : index +// CHECK: %[[A4:.*]] = call @createSparseTensorReader(%[[A0]]) +// CHECK: %[[A5:.*]] = memref.alloca() : memref<2xindex> +// CHECK: %[[A6:.*]] = memref.cast %[[A5]] : memref<2xindex> to memref +// CHECK: call @copySparseTensorReaderDimSizes(%[[A4]], %[[A6]]) +// CHECK: %[[A7:.*]] = memref.load %[[A5]]{{\[}}%[[A2]]] : memref<2xindex> +// CHECK: %[[A8:.*]] = memref.load %[[A5]]{{\[}}%[[A1]]] : memref<2xindex> +// CHECK: %[[A9:.*]] = call @getSparseTensorReaderNNZ(%[[A4]]) +// CHECK: %[[A10:.*]] = arith.muli %[[A9]], %[[A3]] : index +// CHECK: %[[A11:.*]] = memref.alloc() : memref<2xindex> +// CHECK: %[[A12:.*]] = memref.cast %[[A11]] : memref<2xindex> to memref +// CHECK: %[[A13:.*]] = memref.alloc(%[[A10]]) : memref +// CHECK: %[[A14:.*]] = memref.alloc(%[[A9]]) : memref +// CHECK: %[[A15:.*]] = sparse_tensor.storage_specifier.init : !sparse_tensor.storage_specifier<#sparse_tensor.encoding<{ dimLevelType = [ "compressed", "singleton" ] }>> +// CHECK: %[[A16:.*]] = arith.index_cast %[[A8]] : index to i64 +// CHECK: %[[A17:.*]] = sparse_tensor.storage_specifier.set %[[A15]] dim_sz at 0 with %[[A16]] +// CHECK: %[[A18:.*]] = sparse_tensor.storage_specifier.get %[[A17]] ptr_mem_sz at 0 +// CHECK: %[[A19:.*]] = arith.index_cast %[[A18]] : i64 to index +// CHECK: %[[A20:.*]], %[[A21:.*]] = sparse_tensor.push_back %[[A19]], %[[A12]], %[[A2]] +// CHECK: %[[A22:.*]] = arith.index_cast %[[A21]] : index to i64 +// CHECK: %[[A23:.*]] = sparse_tensor.storage_specifier.set %[[A17]] ptr_mem_sz at 0 with %[[A22]] +// CHECK: %[[A24:.*]] = arith.index_cast %[[A7]] : index to i64 +// CHECK: %[[A25:.*]] = sparse_tensor.storage_specifier.set %[[A23]] dim_sz at 1 with %[[A24]] +// CHECK: %[[A26:.*]], %[[A27:.*]] = sparse_tensor.push_back %[[A21]], %[[A20]], %[[A2]], %[[A1]] +// CHECK: %[[A28:.*]] = arith.index_cast %[[A27]] : index to i64 +// CHECK: %[[A29:.*]] = sparse_tensor.storage_specifier.set %[[A25]] ptr_mem_sz at 0 with %[[A28]] +// CHECK: %[[A30:.*]] = memref.alloca() : memref<2xindex> +// CHECK: %[[A31:.*]] = memref.cast %[[A30]] : memref<2xindex> to memref +// CHECK: memref.store %[[A1]], %[[A30]]{{\[}}%[[A2]]] : memref<2xindex> +// CHECK: memref.store %[[A2]], %[[A30]]{{\[}}%[[A1]]] : memref<2xindex> +// CHECK: %[[A32:.*]] = call @getSparseTensorReaderRead0F32(%[[A4]], %[[A31]], %[[A13]], %[[A14]]) +// CHECK: memref.store %[[A9]], %[[A26]]{{\[}}%[[A1]]] : memref +// CHECK: %[[A33:.*]] = arith.index_cast %[[A10]] : index to i64 +// CHECK: %[[A34:.*]] = sparse_tensor.storage_specifier.set %[[A29]] idx_mem_sz at 0 with %[[A33]] +// CHECK: %[[A35:.*]] = arith.index_cast %[[A9]] : index to i64 +// CHECK: %[[A36:.*]] = sparse_tensor.storage_specifier.set %[[A34]] val_mem_sz with %[[A35]] +// CHECK: call @delSparseTensorReader(%[[A4]]) : (!llvm.ptr) -> () +// CHECK: return %[[A26]], %[[A13]], %[[A14]], %[[A36]] +func.func @sparse_new_coo_permute_no(%arg0: !llvm.ptr) -> tensor { + %0 = sparse_tensor.new %arg0 : !llvm.ptr to tensor + return %0 : tensor +} \ No newline at end of file diff --git a/mlir/test/Dialect/SparseTensor/rewriting_for_codegen.mlir b/mlir/test/Dialect/SparseTensor/rewriting_for_codegen.mlir --- a/mlir/test/Dialect/SparseTensor/rewriting_for_codegen.mlir +++ b/mlir/test/Dialect/SparseTensor/rewriting_for_codegen.mlir @@ -10,31 +10,15 @@ dimOrdering = affine_map<(i, j) -> (j, i)> }> +#COO = #sparse_tensor.encoding<{ + dimLevelType = [ "compressed-nu", "singleton" ] +}> + // CHECK-LABEL: func.func @sparse_new( // CHECK-SAME: %[[A:.*]]: !llvm.ptr) -> tensor> { -// CHECK-DAG: %[[C2:.*]] = arith.constant 2 : index -// CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index -// CHECK-DAG: %[[C1:.*]] = arith.constant 1 : index -// CHECK: %[[R:.*]] = call @createSparseTensorReader(%[[A]]) -// CHECK: %[[DS:.*]] = memref.alloca(%[[C2]]) : memref -// CHECK: call @copySparseTensorReaderDimSizes(%[[R]], %[[DS]]) -// CHECK: %[[D0:.*]] = memref.load %[[DS]]{{\[}}%[[C0]]] -// CHECK: %[[D1:.*]] = memref.load %[[DS]]{{\[}}%[[C1]]] -// CHECK: %[[N:.*]] = call @getSparseTensorReaderNNZ(%[[R]]) -// CHECK: %[[T:.*]] = bufferization.alloc_tensor(%[[D0]], %[[D1]]) size_hint=%[[N]] -// CHECK: %[[VB:.*]] = memref.alloca() -// CHECK: %[[T2:.*]] = scf.for %{{.*}} = %[[C0]] to %[[N]] step %[[C1]] iter_args(%[[A2:.*]] = %[[T]]) -// CHECK: func.call @getSparseTensorReaderNextF32(%[[R]], %[[DS]], %[[VB]]) -// CHECK: %[[E0:.*]] = memref.load %[[DS]]{{\[}}%[[C0]]] -// CHECK: %[[E1:.*]] = memref.load %[[DS]]{{\[}}%[[C1]]] -// CHECK: %[[V:.*]] = memref.load %[[VB]][] -// CHECK: %[[T1:.*]] = sparse_tensor.insert %[[V]] into %[[A2]]{{\[}}%[[E0]], %[[E1]]] -// CHECK: scf.yield %[[T1]] -// CHECK: } -// CHECK: call @delSparseTensorReader(%[[R]]) -// CHECK: %[[T4:.*]] = sparse_tensor.load %[[T2]] hasInserts -// CHECK: %[[R:.*]] = sparse_tensor.convert %[[T4]] -// CHECK: bufferization.dealloc_tensor %[[T4]] +// CHECK: %[[COO:.*]] = sparse_tensor.new %[[A]] : !llvm.ptr to tensor> +// CHECK: %[[R:.*]] = sparse_tensor.convert %[[COO]] +// CHECK: bufferization.dealloc_tensor %[[COO]] // CHECK: return %[[R]] func.func @sparse_new(%arg0: !llvm.ptr) -> tensor { %0 = sparse_tensor.new %arg0 : !llvm.ptr to tensor @@ -43,35 +27,24 @@ // CHECK-LABEL: func.func @sparse_new_csc( // CHECK-SAME: %[[A:.*]]: !llvm.ptr) -> tensor (d1, d0)> }>> { -// CHECK-DAG: %[[C2:.*]] = arith.constant 2 : index -// CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index -// CHECK-DAG: %[[C1:.*]] = arith.constant 1 : index -// CHECK: %[[R:.*]] = call @createSparseTensorReader(%[[A]]) -// CHECK: %[[DS:.*]] = memref.alloca(%[[C2]]) : memref -// CHECK: call @copySparseTensorReaderDimSizes(%[[R]], %[[DS]]) -// CHECK: %[[D0:.*]] = memref.load %[[DS]]{{\[}}%[[C0]]] -// CHECK: %[[D1:.*]] = memref.load %[[DS]]{{\[}}%[[C1]]] -// CHECK: %[[N:.*]] = call @getSparseTensorReaderNNZ(%[[R]]) -// CHECK: %[[T:.*]] = bufferization.alloc_tensor(%[[D0]], %[[D1]]) size_hint=%[[N]] -// CHECK: %[[VB:.*]] = memref.alloca() -// CHECK: %[[T2:.*]] = scf.for %{{.*}} = %[[C0]] to %[[N]] step %[[C1]] iter_args(%[[A2:.*]] = %[[T]]) -// CHECK: func.call @getSparseTensorReaderNextF32(%[[R]], %[[DS]], %[[VB]]) -// CHECK: %[[E0:.*]] = memref.load %[[DS]]{{\[}}%[[C0]]] -// CHECK: %[[E1:.*]] = memref.load %[[DS]]{{\[}}%[[C1]]] -// CHECK: %[[V:.*]] = memref.load %[[VB]][] -// CHECK: %[[T1:.*]] = sparse_tensor.insert %[[V]] into %[[A2]]{{\[}}%[[E1]], %[[E0]]] -// CHECK: scf.yield %[[T1]] -// CHECK: } -// CHECK: call @delSparseTensorReader(%[[R]]) -// CHECK: %[[T4:.*]] = sparse_tensor.load %[[T2]] hasInserts -// CHECK: %[[R:.*]] = sparse_tensor.convert %[[T4]] -// CHECK: bufferization.dealloc_tensor %[[T4]] +// CHECK: %[[COO:.*]] = sparse_tensor.new %[[A]] : !llvm.ptr to tensor (d1, d0)> }>> +// CHECK: %[[R:.*]] = sparse_tensor.convert %[[COO]] +// CHECK: bufferization.dealloc_tensor %[[COO]] // CHECK: return %[[R]] func.func @sparse_new_csc(%arg0: !llvm.ptr) -> tensor { %0 = sparse_tensor.new %arg0 : !llvm.ptr to tensor return %0 : tensor } +// CHECK-LABEL: func.func @sparse_new_coo( +// CHECK-SAME: %[[A:.*]]: !llvm.ptr) -> tensor> { +// CHECK: %[[COO:.*]] = sparse_tensor.new %[[A]] : !llvm.ptr to tensor> +// CHECK: return %[[COO]] +func.func @sparse_new_coo(%arg0: !llvm.ptr) -> tensor { + %0 = sparse_tensor.new %arg0 : !llvm.ptr to tensor + return %0 : tensor +} + // CHECK-LABEL: func.func @sparse_out( // CHECK-SAME: %[[A:.*]]: tensor<10x20xf32, #sparse_tensor.encoding<{ dimLevelType = [ "dense", "compressed" ] }>>, // CHECK-SAME: %[[B:.*]]: !llvm.ptr) {