[Mlir-commits] [mlir] 2c81d43 - [mlir][sparse] Improve the implementation of sparse_tensor.new for the codegen path.
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Wed Mar 1 07:29:55 PST 2023
Author: bixia1
Date: 2023-03-01T07:29:49-08:00
New Revision: 2c81d43241f79168dc3f63691c1f154d33942dc9
URL: https://github.com/llvm/llvm-project/commit/2c81d43241f79168dc3f63691c1f154d33942dc9
DIFF: https://github.com/llvm/llvm-project/commit/2c81d43241f79168dc3f63691c1f154d33942dc9.diff
LOG: [mlir][sparse] Improve the implementation of sparse_tensor.new for the codegen path.
Rewrite a NewOp into a NewOp of a sorted COO tensor and a ConvertOp for
converting the sorted COO tensor to the destination tensor type.
Codegen a NewOp of a sorted COO tensor to use the new bulk reader API and sort
the elements only when the input is not sorted.
Reviewed By: aartbik
Differential Revision: https://reviews.llvm.org/D144504
Added:
Modified:
mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorCodegen.cpp
mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorRewriting.cpp
mlir/test/Dialect/SparseTensor/codegen.mlir
mlir/test/Dialect/SparseTensor/rewriting_for_codegen.mlir
Removed:
################################################################################
diff --git a/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorCodegen.cpp b/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorCodegen.cpp
index f3f67846877aa..7da6e0f996f92 100644
--- a/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorCodegen.cpp
+++ b/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorCodegen.cpp
@@ -1289,6 +1289,137 @@ struct SparseUnpackOpConverter : public OpConversionPattern<UnpackOp> {
}
};
+struct SparseNewOpConverter : public OpConversionPattern<NewOp> {
+ 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();
+ // Creating COO with NewOp is handled by direct IR codegen. All other cases
+ // are handled by rewriting.
+ if (!dstTp.hasEncoding() || getCOOStart(encDst) != 0)
+ return failure();
+
+ // Implement the NewOp(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<Value> dynSizes;
+ if (dstTp.hasDynamicDimShape()) {
+ Value dimSizes = genAlloca(rewriter, loc, dimRank, indexTp);
+ createFuncCall(rewriter, loc, "copySparseTensorReaderDimSizes", {},
+ {reader, dimSizes}, EmitCInterface::On)
+ .getResult(0);
+ ArrayRef<int64_t> dstShape = dstTp.getRankedTensorType().getShape();
+ for (auto &d : llvm::enumerate(dstShape)) {
+ if (d.value() == ShapedType::kDynamic) {
+ dynSizes.push_back(rewriter.create<memref::LoadOp>(
+ 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<Value> 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> getReadFuncName{"getSparseTensorReaderRead",
+ overheadTypeFunctionSuffix(indBufEleTp),
+ primaryTypeFunctionSuffix(eltTp)};
+
+ Value xs = desc.getAOSMemRef();
+ Value ys = desc.getValMemRef();
+ SmallVector<Value> 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, getReadFuncName, {f.getType()},
+ {reader, dim2lvl, xs, ys}, EmitCInterface::On)
+ .getResult(0);
+
+ // If the destination tensor is a sorted COO, we need to sort the COO tensor
+ // data if the input elements aren't sorted yet.
+ if (encDst.isOrderedLvl(dimRank - 1)) {
+ Value notSorted = rewriter.create<arith::CmpIOp>(
+ loc, arith::CmpIPredicate::eq, isSorted, f);
+ scf::IfOp ifOp =
+ rewriter.create<scf::IfOp>(loc, notSorted, /*else*/ false);
+ rewriter.setInsertionPointToStart(&ifOp.getThenRegion().front());
+ rewriter.create<SortCooOp>(
+ 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<arith::IndexCastOp>(loc, ptrEleTy, nse);
+ rewriter.create<memref::StoreOp>(loc, ptrNse, ptrMemref0, c1);
+
+ // Update storage specifier.
+ Value idxSize = rewriter.create<arith::MulIOp>(
+ 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
//===----------------------------------------------------------------------===//
@@ -1308,8 +1439,8 @@ void mlir::populateSparseTensorCodegenPatterns(
SparseInsertConverter, SparseToPointersConverter,
SparseToIndicesConverter, SparseToIndicesBufferConverter,
SparseToValuesConverter, SparseConvertConverter,
- SparseNumberOfEntriesConverter>(typeConverter,
- patterns.getContext());
+ SparseNewOpConverter, SparseNumberOfEntriesConverter>(
+ typeConverter, patterns.getContext());
patterns.add<SparseTensorAllocConverter>(typeConverter, patterns.getContext(),
enableBufferInitialization);
}
diff --git a/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorRewriting.cpp b/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorRewriting.cpp
index 20713609c847f..b128669a707f8 100644
--- a/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorRewriting.cpp
+++ b/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorRewriting.cpp
@@ -1013,92 +1013,21 @@ struct NewRewriter : public OpRewritePattern<NewOp> {
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<Value> dynSizesArray;
- if (dstTp.hasDynamicDimShape()) {
- createFuncCall(rewriter, loc, "copySparseTensorReaderDimSizes", {},
- {reader, dimSizes}, EmitCInterface::On)
- .getResult(0);
- ArrayRef<int64_t> dstShape = dstTp.getRankedTensorType().getShape();
- for (auto &d : llvm::enumerate(dstShape)) {
- if (d.value() == ShapedType::kDynamic) {
- dynSizesArray.push_back(rewriter.create<memref::LoadOp>(
- 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<AllocTensorOp>(loc, cooTp, dynSizesArray, Value(),
- /*sizeHint=*/nnz, Attribute())
- .getResult();
-
- Type eltTp = dstTp.getElementType();
- Value value = genAllocaScalar(rewriter, loc, eltTp);
- scf::ForOp forOp = rewriter.create<scf::ForOp>(loc, c0, nnz, c1,
- ArrayRef<Value>(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<Value> indicesArray(dimRank);
- for (Dimension d = 0; d < dimRank; d++) {
- // FIXME: `toStoredDim` is deprecated
- indicesArray[toStoredDim(encDst, d)] = rewriter.create<memref::LoadOp>(
- loc, indices, constantIndex(rewriter, loc, d));
- }
- Value v = rewriter.create<memref::LoadOp>(loc, value);
- Value t = rewriter.create<InsertOp>(loc, v, forOp.getRegionIterArg(0),
- indicesArray);
- rewriter.create<scf::YieldOp>(loc, ArrayRef<Value>(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<LoadOp>(loc, cooBuffer, true);
- Value newOp = rewriter.replaceOpWithNewOp<ConvertOp>(
- op, dstTp.getRankedTensorType(), cooBuffer);
-
- // Release the unordered COO tensor buffer.
- rewriter.setInsertionPointAfterValue(newOp);
- rewriter.create<DeallocTensorOp>(loc, cooBuffer);
+ getCOOFromTypeWithOrdering(dstTp, encDst.getDimOrdering(), true);
+ Value cooTensor = rewriter.create<NewOp>(loc, cooTp, op.getSource());
+ Value convert = rewriter.replaceOpWithNewOp<ConvertOp>(
+ op, dstTp.getRankedTensorType(), cooTensor);
+
+ // Release the ordered COO tensor.
+ rewriter.setInsertionPointAfterValue(convert);
+ rewriter.create<DeallocTensorOp>(loc, cooTensor);
return success();
}
diff --git a/mlir/test/Dialect/SparseTensor/codegen.mlir b/mlir/test/Dialect/SparseTensor/codegen.mlir
index 4b1beb51713e7..827637bb4bc3b 100644
--- 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" ]
}>
@@ -657,3 +662,88 @@ func.func @sparse_nop_convert(%arg0: tensor<32xf32, #SparseVector>) -> tensor<?x
%0 = sparse_tensor.convert %arg0 : tensor<32xf32, #SparseVector> to tensor<?xf32, #SparseVector>
return %0 : tensor<?xf32, #SparseVector>
}
+
+// CHECK-LABEL: func.func @sparse_new_coo(
+// CHECK-SAME: %[[A0:.*]]: !llvm.ptr<i8>) -> (memref<?xindex>, memref<?xindex>, memref<?xf32>, !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<?xindex>
+// CHECK: call @copySparseTensorReaderDimSizes(%[[A5]], %[[A7]]) : (!llvm.ptr<i8>, memref<?xindex>) -> ()
+// 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<?xindex>
+// CHECK: %[[A14:.*]] = memref.alloc(%[[A11]]) : memref<?xindex>
+// CHECK: %[[A15:.*]] = memref.alloc(%[[A10]]) : memref<?xf32>
+// CHECK: %[[A16:.*]] = sparse_tensor.storage_specifier.init : !sparse_tensor.storage_specifier<#sparse_tensor.encoding<{ dimLevelType = [ "compressed", "singleton" ] }>>
+// CHECK: %[[A18:.*]] = sparse_tensor.storage_specifier.set %[[A16]] dim_sz at 0 with %[[A8]]
+// CHECK: %[[A19:.*]] = sparse_tensor.storage_specifier.get %[[A18]] ptr_mem_sz at 0
+// CHECK: %[[A21:.*]], %[[A22:.*]] = sparse_tensor.push_back %[[A19]], %[[A13]], %[[A3]]
+// CHECK: %[[A24:.*]] = sparse_tensor.storage_specifier.set %[[A18]] ptr_mem_sz at 0 with %[[A22]]
+// CHECK: %[[A26:.*]] = sparse_tensor.storage_specifier.set %[[A24]] dim_sz at 1 with %[[A9]]
+// CHECK: %[[A27:.*]], %[[A28:.*]] = sparse_tensor.push_back %[[A22]], %[[A21]], %[[A3]], %[[A2]]
+// CHECK: %[[A30:.*]] = sparse_tensor.storage_specifier.set %[[A26]] ptr_mem_sz at 0 with %[[A28]]
+// CHECK: %[[A31:.*]] = memref.alloca() : memref<2xindex>
+// CHECK: %[[A32:.*]] = memref.cast %[[A31]] : memref<2xindex> to memref<?xindex>
+// 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<?xindex> jointly memref<?xf32>
+// CHECK: }
+// CHECK: memref.store %[[A10]], %[[A27]]{{\[}}%[[A2]]] : memref<?xindex>
+// CHECK: %[[A36:.*]] = sparse_tensor.storage_specifier.set %[[A30]] idx_mem_sz at 0 with %[[A11]]
+// CHECK: %[[A38:.*]] = sparse_tensor.storage_specifier.set %[[A36]] val_mem_sz with %[[A10]]
+// CHECK: call @delSparseTensorReader(%[[A5]]) : (!llvm.ptr<i8>) -> ()
+// CHECK: return %[[A27]], %[[A14]], %[[A15]], %[[A38]]
+func.func @sparse_new_coo(%arg0: !llvm.ptr<i8>) -> tensor<?x?xf32, #Coo> {
+ %0 = sparse_tensor.new %arg0 : !llvm.ptr<i8> to tensor<?x?xf32, #Coo>
+ return %0 : tensor<?x?xf32, #Coo>
+}
+
+// CHECK-LABEL: func.func @sparse_new_coo_permute_no(
+// CHECK-SAME: %[[A0:.*]]: !llvm.ptr<i8>) -> (memref<?xindex>, memref<?xindex>, memref<?xf32>, !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<?xindex>
+// 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<?xindex>
+// CHECK: %[[A13:.*]] = memref.alloc(%[[A10]]) : memref<?xindex>
+// CHECK: %[[A14:.*]] = memref.alloc(%[[A9]]) : memref<?xf32>
+// CHECK: %[[A15:.*]] = sparse_tensor.storage_specifier.init : !sparse_tensor.storage_specifier<#sparse_tensor.encoding<{ dimLevelType = [ "compressed", "singleton" ] }>>
+// CHECK: %[[A17:.*]] = sparse_tensor.storage_specifier.set %[[A15]] dim_sz at 0 with %[[A8]]
+// CHECK: %[[A18:.*]] = sparse_tensor.storage_specifier.get %[[A17]] ptr_mem_sz at 0
+// CHECK: %[[A20:.*]], %[[A21:.*]] = sparse_tensor.push_back %[[A18]], %[[A12]], %[[A2]]
+// CHECK: %[[A23:.*]] = sparse_tensor.storage_specifier.set %[[A17]] ptr_mem_sz at 0 with %[[A21]]
+// CHECK: %[[A25:.*]] = sparse_tensor.storage_specifier.set %[[A23]] dim_sz at 1 with %[[A7]]
+// CHECK: %[[A26:.*]], %[[A27:.*]] = sparse_tensor.push_back %[[A21]], %[[A20]], %[[A2]], %[[A1]]
+// CHECK: %[[A29:.*]] = sparse_tensor.storage_specifier.set %[[A25]] ptr_mem_sz at 0 with %[[A27]]
+// CHECK: %[[A30:.*]] = memref.alloca() : memref<2xindex>
+// CHECK: %[[A31:.*]] = memref.cast %[[A30]] : memref<2xindex> to memref<?xindex>
+// 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<?xindex>
+// CHECK: %[[A34:.*]] = sparse_tensor.storage_specifier.set %[[A29]] idx_mem_sz at 0 with %[[A10]]
+// CHECK: %[[A36:.*]] = sparse_tensor.storage_specifier.set %[[A34]] val_mem_sz with %[[A9]]
+// CHECK: call @delSparseTensorReader(%[[A4]]) : (!llvm.ptr<i8>) -> ()
+// CHECK: return %[[A26]], %[[A13]], %[[A14]], %[[A36]]
+func.func @sparse_new_coo_permute_no(%arg0: !llvm.ptr<i8>) -> tensor<?x?xf32, #CooPNo> {
+ %0 = sparse_tensor.new %arg0 : !llvm.ptr<i8> to tensor<?x?xf32, #CooPNo>
+ return %0 : tensor<?x?xf32, #CooPNo>
+}
\ 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
index a3ce325e8fbef..db2504597bb55 100644
--- 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<i8>) -> tensor<?x?xf32, #sparse_tensor.encoding<{ dimLevelType = [ "dense", "compressed" ] }>> {
-// 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<?xindex>
-// 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<i8> to tensor<?x?xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed-nu", "singleton" ] }>>
+// CHECK: %[[R:.*]] = sparse_tensor.convert %[[COO]]
+// CHECK: bufferization.dealloc_tensor %[[COO]]
// CHECK: return %[[R]]
func.func @sparse_new(%arg0: !llvm.ptr<i8>) -> tensor<?x?xf32, #CSR> {
%0 = sparse_tensor.new %arg0 : !llvm.ptr<i8> to tensor<?x?xf32, #CSR>
@@ -43,35 +27,24 @@ func.func @sparse_new(%arg0: !llvm.ptr<i8>) -> tensor<?x?xf32, #CSR> {
// CHECK-LABEL: func.func @sparse_new_csc(
// CHECK-SAME: %[[A:.*]]: !llvm.ptr<i8>) -> tensor<?x?xf32, #sparse_tensor.encoding<{ dimLevelType = [ "dense", "compressed" ], dimOrdering = affine_map<(d0, d1) -> (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<?xindex>
-// 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<i8> to tensor<?x?xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed-nu", "singleton" ], dimOrdering = affine_map<(d0, d1) -> (d1, d0)> }>>
+// CHECK: %[[R:.*]] = sparse_tensor.convert %[[COO]]
+// CHECK: bufferization.dealloc_tensor %[[COO]]
// CHECK: return %[[R]]
func.func @sparse_new_csc(%arg0: !llvm.ptr<i8>) -> tensor<?x?xf32, #CSC> {
%0 = sparse_tensor.new %arg0 : !llvm.ptr<i8> to tensor<?x?xf32, #CSC>
return %0 : tensor<?x?xf32, #CSC>
}
+// CHECK-LABEL: func.func @sparse_new_coo(
+// CHECK-SAME: %[[A:.*]]: !llvm.ptr<i8>) -> tensor<?x?xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed-nu", "singleton" ] }>> {
+// CHECK: %[[COO:.*]] = sparse_tensor.new %[[A]] : !llvm.ptr<i8> to tensor<?x?xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed-nu", "singleton" ] }>>
+// CHECK: return %[[COO]]
+func.func @sparse_new_coo(%arg0: !llvm.ptr<i8>) -> tensor<?x?xf32, #COO> {
+ %0 = sparse_tensor.new %arg0 : !llvm.ptr<i8> to tensor<?x?xf32, #COO>
+ return %0 : tensor<?x?xf32, #COO>
+}
+
// CHECK-LABEL: func.func @sparse_out(
// CHECK-SAME: %[[A:.*]]: tensor<10x20xf32, #sparse_tensor.encoding<{ dimLevelType = [ "dense", "compressed" ] }>>,
// CHECK-SAME: %[[B:.*]]: !llvm.ptr<i8>) {
More information about the Mlir-commits
mailing list