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 @@ -80,9 +80,9 @@ string summary = "Converts between different tensor types"; string description = [{ Converts one sparse or dense tensor type to another tensor type. The rank - and dimensions of the source and destination types must match exactly, - only the sparse encoding of these types may be different. The name `convert` - was preferred over `cast`, since the operation may incur a non-trivial cost. + and dimensions of the source and destination types must match, but the sparse + encoding of these types can obviously be different. The name `convert` was + preferred over `cast`, since the operation may incur a non-trivial cost. When converting between two different sparse tensor types, only explicitly stored values are moved from one underlying sparse storage format to @@ -97,9 +97,14 @@ Examples: ```mlir - %0 = sparse_tensor.convert %1 : tensor<32x32xf32> to tensor<32x32xf32, #CSR> - - %2 = sparse_tensor.convert %3 : tensor<8x8xi32, #CSC> to tensor<8x8xi32, #CSR> + %0 = sparse_tensor.convert %a : tensor<32x32xf32> to tensor<32x32xf32, #CSR> + %1 = sparse_tensor.convert %a : tensor<32x32xf32> to tensor + %2 = sparse_tensor.convert %b : tensor<8x8xi32, #CSC> to tensor<8x8xi32, #CSR> + %3 = sparse_tensor.convert %c : tensor<4x8xf64, #CSR> to tensor<4x?xf64, #CSC> + + // The following conversion is not allowed (since it would require a + // runtime assertion that the source's dimension size is actually 100). + %4 = sparse_tensor.convert %d : tensor to tensor<100xf64, #SV> ``` }]; 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 @@ -240,8 +240,11 @@ assert(tp1.getRank() == tp2.getRank()); auto shape1 = tp1.getShape(); auto shape2 = tp2.getShape(); + // Accept size matches between the source and the destination type + // (e.g. 10 vs. 10, 10 vs. ?, or ? vs. ?), but reject direct mismatches or + // matches that would need a runtime assert (e.g. 10 vs. 20 or ? vs. 10). for (unsigned d = 0, rank = tp1.getRank(); d < rank; d++) { - if (shape1[d] != shape2[d]) + if (shape1[d] != shape2[d] && shape2[d] != ShapedType::kDynamicSize) return op.emitError("unexpected conversion mismatch in dimension ") << d; } 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 @@ -99,7 +99,7 @@ /// Generates a constant of `index` type. inline static Value constantIndex(ConversionPatternRewriter &rewriter, - Location loc, unsigned i) { + Location loc, int64_t i) { return rewriter.create(loc, i); } @@ -144,6 +144,70 @@ return result; } +/// Generates dimension size call. +static Value genDimSizeCall(ConversionPatternRewriter &rewriter, Operation *op, + SparseTensorEncodingAttr &enc, Value src, + int64_t idx) { + // Permute the index according to an optional dimension ordering. + if (AffineMap p = enc.getDimOrdering()) + idx = p.getPermutedPosition(idx); + // Generate the call. + Location loc = op->getLoc(); + StringRef name = "sparseDimSize"; + SmallVector params; + params.push_back(src); + params.push_back(constantIndex(rewriter, loc, idx)); + Type iTp = rewriter.getIndexType(); + auto fn = getFunc(op, name, iTp, params); + return rewriter.create(loc, iTp, fn, params).getResult(0); +} + +/// Generates a call into the "swiss army knife" method of the sparse runtime +/// support library for materializing sparse tensors into the computation. +static Value genNewCall(ConversionPatternRewriter &rewriter, Operation *op, + ArrayRef params) { + Location loc = op->getLoc(); + StringRef name = "newSparseTensor"; + Type pTp = LLVM::LLVMPointerType::get(rewriter.getI8Type()); + auto fn = getFunc(op, name, pTp, params, /*emitCInterface=*/true); + auto call = rewriter.create(loc, pTp, fn, params); + return call.getResult(0); +} + +/// Populates given sizes array from type. +static void sizesFromType(ConversionPatternRewriter &rewriter, + SmallVector &sizes, Location loc, + ShapedType stp) { + auto shape = stp.getShape(); + for (unsigned i = 0, rank = stp.getRank(); i < rank; i++) { + uint64_t s = shape[i] == ShapedType::kDynamicSize ? 0 : shape[i]; + sizes.push_back(constantIndex(rewriter, loc, s)); + } +} + +/// Populates given sizes array from source. +static void sizesFromSrc(ConversionPatternRewriter &rewriter, + SmallVector &sizes, Location loc, + Value src) { + ShapedType stp = src.getType().cast(); + for (unsigned i = 0, rank = stp.getRank(); i < rank; i++) + sizes.push_back(linalg::createOrFoldDimOp(rewriter, loc, src, i)); +} + +/// Populates given sizes array from type (for static sizes) and from +/// an already converted into opague pointer source (for dynamic sizes). +static void sizesFromPtr(ConversionPatternRewriter &rewriter, + SmallVector &sizes, Operation *op, + SparseTensorEncodingAttr &enc, ShapedType stp, + Value src) { + auto shape = stp.getShape(); + for (unsigned i = 0, rank = stp.getRank(); i < rank; i++) + if (shape[i] == ShapedType::kDynamicSize) + sizes.push_back(genDimSizeCall(rewriter, op, enc, src, i)); + else + sizes.push_back(constantIndex(rewriter, op->getLoc(), shape[i])); +} + /// Generates a temporary buffer of the given size and type. static Value genAlloca(ConversionPatternRewriter &rewriter, Location loc, unsigned sz, Type tp) { @@ -152,7 +216,7 @@ return rewriter.create(loc, memTp, ValueRange{a}); } -/// Fills a temporary buffer of the given type with arguments. +/// Generates a temporary buffer of the given type and given contents. static Value genBuffer(ConversionPatternRewriter &rewriter, Location loc, ArrayRef values) { unsigned sz = values.size(); @@ -165,36 +229,28 @@ return buffer; } -/// Generates a call into the "swiss army knife" method of the sparse runtime -/// support library for materializing sparse tensors into the computation. The -/// method returns the call value and assigns the permutation to 'perm'. -static Value genNewCall(ConversionPatternRewriter &rewriter, Operation *op, - SparseTensorEncodingAttr &enc, uint32_t action, - Value &perm, ValueRange szs, Value ptr = Value()) { +/// Populates parameters required to call the "swiss army knife" method of the +/// sparse runtime support library for materializing sparse tensors into the +/// computation. +static void newParams(ConversionPatternRewriter &rewriter, + SmallVector ¶ms, Operation *op, + SparseTensorEncodingAttr &enc, uint32_t action, + ValueRange szs, Value ptr = Value()) { Location loc = op->getLoc(); - ShapedType resType = op->getResult(0).getType().cast(); - SmallVector params; - // Sparsity annotations in tensor constant form. - SmallVector attrs; ArrayRef dlt = enc.getDimLevelType(); unsigned sz = dlt.size(); + // Sparsity annotations. + SmallVector attrs; for (unsigned i = 0; i < sz; i++) attrs.push_back(constantI8(rewriter, loc, getDimLevelTypeEncoding(dlt[i]))); params.push_back(genBuffer(rewriter, loc, attrs)); - // Dimension sizes array of the enveloping *dense* tensor. Useful for either + // Dimension sizes array of the enveloping tensor. Useful for either // verification of external data, or for construction of internal data. - auto shape = resType.getShape(); + // The index type is casted to I64 for API consistency. + Type iTp = rewriter.getI64Type(); SmallVector sizes; - if (szs.size() > 0) { - for (Value s : szs) - sizes.push_back( - rewriter.create(loc, s, rewriter.getI64Type())); - } else { - for (unsigned i = 0; i < sz; i++) { - uint64_t s = shape[i] == ShapedType::kDynamicSize ? 0 : shape[i]; - sizes.push_back(constantI64(rewriter, loc, s)); - } - } + for (Value s : szs) + sizes.push_back(rewriter.create(loc, s, iTp)); params.push_back(genBuffer(rewriter, loc, sizes)); // Dimension order permutation array. This is the "identity" permutation by // default, or otherwise the "reverse" permutation of a given ordering, so @@ -207,9 +263,9 @@ for (unsigned i = 0; i < sz; i++) rev[i] = constantI64(rewriter, loc, i); } - perm = genBuffer(rewriter, loc, rev); - params.push_back(perm); + params.push_back(genBuffer(rewriter, loc, rev)); // Secondary and primary types encoding. + ShapedType resType = op->getResult(0).getType().cast(); unsigned secPtr = getOverheadTypeEncoding(enc.getPointerBitWidth()); unsigned secInd = getOverheadTypeEncoding(enc.getIndexBitWidth()); unsigned primary = getPrimaryTypeEncoding(resType.getElementType()); @@ -223,12 +279,6 @@ ptr = rewriter.create(loc, pTp); params.push_back(constantI32(rewriter, loc, action)); params.push_back(ptr); - // Generate the call to create new tensor. - StringRef name = "newSparseTensor"; - auto call = rewriter.create( - loc, pTp, getFunc(op, name, pTp, params, /*emitCInterface=*/true), - params); - return call.getResult(0); } /// Generates the comparison `v != 0` where `v` is of numeric type `t`. @@ -299,9 +349,8 @@ params.push_back(ind); params.push_back(perm); Type pTp = LLVM::LLVMPointerType::get(rewriter.getI8Type()); - rewriter.create( - loc, pTp, getFunc(op, name, pTp, params, /*emitCInterface=*/true), - params); + auto fn = getFunc(op, name, pTp, params, /*emitCInterface=*/true); + rewriter.create(loc, pTp, fn, params); } /// If the tensor is a sparse constant, generates and returns the pair of @@ -362,24 +411,17 @@ LogicalResult matchAndRewrite(tensor::DimOp op, OpAdaptor adaptor, ConversionPatternRewriter &rewriter) const override { - Type resType = op.getType(); + // Only rewrite annotated DimOp with constant index. auto enc = getSparseTensorEncoding(op.source().getType()); if (!enc) return failure(); - // Permute the dim index. Optional index = op.getConstantIndex(); if (!index.hasValue()) return failure(); - int64_t idx = index.getValue(); - if (AffineMap p = enc.getDimOrdering()) - idx = p.getPermutedPosition(idx); // Generate the call. - StringRef name = "sparseDimSize"; - SmallVector params; - params.push_back(adaptor.getOperands()[0]); - params.push_back(constantIndex(rewriter, op.getLoc(), idx)); - rewriter.replaceOpWithNewOp( - op, resType, getFunc(op, name, resType, params), params); + Value src = adaptor.getOperands()[0]; + int64_t idx = index.getValue(); + rewriter.replaceOp(op, genDimSizeCall(rewriter, op, enc, src, idx)); return success(); } }; @@ -394,9 +436,14 @@ auto enc = getSparseTensorEncoding(resType); if (!enc) return failure(); - Value perm; - rewriter.replaceOp(op, genNewCall(rewriter, op, enc, kFromFile, perm, {}, - adaptor.getOperands()[0])); + // Generate the call to construct tensor from ptr. The sizes are + // inferred from the result type of the new operator. + SmallVector sizes; + SmallVector params; + sizesFromType(rewriter, sizes, op.getLoc(), resType.cast()); + Value ptr = adaptor.getOperands()[0]; + newParams(rewriter, params, op, enc, kFromFile, sizes, ptr); + rewriter.replaceOp(op, genNewCall(rewriter, op, params)); return success(); } }; @@ -411,9 +458,11 @@ auto enc = getSparseTensorEncoding(resType); if (!enc) return failure(); - Value perm; - rewriter.replaceOp( - op, genNewCall(rewriter, op, enc, kEmpty, perm, adaptor.getOperands())); + // Generate the call to construct empty tensor. The sizes are + // explicitly defined by the arguments to the init operator. + SmallVector params; + newParams(rewriter, params, op, enc, kEmpty, adaptor.getOperands()); + rewriter.replaceOp(op, genNewCall(rewriter, op, params)); return success(); } }; @@ -424,10 +473,12 @@ LogicalResult matchAndRewrite(ConvertOp op, OpAdaptor adaptor, ConversionPatternRewriter &rewriter) const override { + Location loc = op->getLoc(); Type resType = op.getType(); + Type srcType = op.source().getType(); auto encDst = getSparseTensorEncoding(resType); - auto encSrc = getSparseTensorEncoding(op.source().getType()); - auto src = adaptor.getOperands()[0]; + auto encSrc = getSparseTensorEncoding(srcType); + Value src = adaptor.getOperands()[0]; if (encDst && encSrc) { // This is a sparse => sparse conversion, which is handled as follows: // t = src->toCOO(); ; src to COO in dst order @@ -435,10 +486,15 @@ // Using the coordinate scheme as an intermediate does not always // yield the fastest conversion but avoids the need for a full // O(N^2) conversion matrix. - Value perm; - Value coo = genNewCall(rewriter, op, encDst, kToCOO, perm, {}, src); - rewriter.replaceOp( - op, genNewCall(rewriter, op, encDst, kFromCOO, perm, {}, coo)); + SmallVector sizes; + SmallVector params; + sizesFromPtr(rewriter, sizes, op, encSrc, srcType.cast(), + src); + newParams(rewriter, params, op, encDst, kToCOO, sizes, src); + Value coo = genNewCall(rewriter, op, params); + params[6] = constantI32(rewriter, loc, kFromCOO); + params[7] = coo; + rewriter.replaceOp(op, genNewCall(rewriter, op, params)); return success(); } if (!encDst || encSrc) { @@ -471,12 +527,15 @@ // Also note that the code below only generates the "new" ops and // the loop-nest per se; whereas the entire body of the innermost // loop is generated by genAddElt(). - Location loc = op->getLoc(); - ShapedType shape = resType.cast(); - Value perm; - Value ptr = genNewCall(rewriter, op, encDst, kEmptyCOO, perm, {}); - Value ind = - genAlloca(rewriter, loc, shape.getRank(), rewriter.getIndexType()); + ShapedType stp = resType.cast(); + unsigned rank = stp.getRank(); + SmallVector sizes; + SmallVector params; + sizesFromSrc(rewriter, sizes, loc, src); + newParams(rewriter, params, op, encDst, kEmptyCOO, sizes); + Value ptr = genNewCall(rewriter, op, params); + Value ind = genAlloca(rewriter, loc, rank, rewriter.getIndexType()); + Value perm = params[2]; SmallVector lo; SmallVector hi; SmallVector st; @@ -493,14 +552,13 @@ hi.push_back(linalg::createOrFoldDimOp(rewriter, loc, values, 0)); st.push_back(one); } else { - for (unsigned i = 0, rank = shape.getRank(); i < rank; i++) { + for (unsigned i = 0; i < rank; i++) { lo.push_back(zero); hi.push_back(linalg::createOrFoldDimOp(rewriter, loc, src, i)); st.push_back(one); } } - Type eltType = shape.getElementType(); - unsigned rank = shape.getRank(); + Type eltType = stp.getElementType(); scf::buildLoopNest( rewriter, op.getLoc(), lo, hi, st, {}, [&](OpBuilder &builder, Location loc, ValueRange ivs, @@ -514,8 +572,10 @@ genAddEltCall(rewriter, op, eltType, ptr, val, ind, perm); return {}; }); - rewriter.replaceOp( - op, genNewCall(rewriter, op, encDst, kFromCOO, perm, {}, ptr)); + // Final call to construct sparse tensor storage. + params[6] = constantI32(rewriter, loc, kFromCOO); + params[7] = ptr; + rewriter.replaceOp(op, genNewCall(rewriter, op, params)); return success(); } }; @@ -529,9 +589,8 @@ ConversionPatternRewriter &rewriter) const override { StringRef name = "delSparseTensor"; TypeRange none; - rewriter.create(op.getLoc(), none, - getFunc(op, name, none, adaptor.getOperands()), - adaptor.getOperands()); + auto fn = getFunc(op, name, none, adaptor.getOperands()); + rewriter.create(op.getLoc(), none, fn, adaptor.getOperands()); rewriter.eraseOp(op); return success(); } @@ -560,11 +619,9 @@ name = "sparsePointers8"; else return failure(); - rewriter.replaceOpWithNewOp(op, resType, - getFunc(op, name, resType, - adaptor.getOperands(), - /*emitCInterface=*/true), - adaptor.getOperands()); + auto fn = getFunc(op, name, resType, adaptor.getOperands(), + /*emitCInterface=*/true); + rewriter.replaceOpWithNewOp(op, resType, fn, adaptor.getOperands()); return success(); } }; @@ -591,11 +648,9 @@ name = "sparseIndices8"; else return failure(); - rewriter.replaceOpWithNewOp(op, resType, - getFunc(op, name, resType, - adaptor.getOperands(), - /*emitCInterface=*/true), - adaptor.getOperands()); + auto fn = getFunc(op, name, resType, adaptor.getOperands(), + /*emitCInterface=*/true); + rewriter.replaceOpWithNewOp(op, resType, fn, adaptor.getOperands()); return success(); } }; @@ -624,11 +679,9 @@ name = "sparseValuesI8"; else return failure(); - rewriter.replaceOpWithNewOp(op, resType, - getFunc(op, name, resType, - adaptor.getOperands(), - /*emitCInterface=*/true), - adaptor.getOperands()); + auto fn = getFunc(op, name, resType, adaptor.getOperands(), + /*emitCInterface=*/true); + rewriter.replaceOpWithNewOp(op, resType, fn, adaptor.getOperands()); return success(); } }; 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 @@ -127,8 +127,8 @@ // CHECK-DAG: %[[JJ:.*]] = arith.index_cast %[[J]] : index to i64 // CHECK-DAG: memref.store %[[II]], %[[Q]][%[[C0]]] : memref<2xi64> // CHECK-DAG: memref.store %[[JJ]], %[[Q]][%[[C1]]] : memref<2xi64> -// CHECK: %[[A:.*]] = llvm.mlir.null : !llvm.ptr -// CHECK: %[[T:.*]] = call @newSparseTensor(%[[X]], %[[Y]], %[[Z]], %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %[[A]]) +// CHECK: %[[NP:.*]] = llvm.mlir.null : !llvm.ptr +// CHECK: %[[T:.*]] = call @newSparseTensor(%[[X]], %[[Y]], %[[Z]], %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %[[NP]]) // CHECK: return %[[T]] : !llvm.ptr func @sparse_init(%arg0: index, %arg1: index) -> tensor { %0 = sparse_tensor.init [%arg0, %arg1] : tensor @@ -156,22 +156,23 @@ // CHECK-SAME: %[[A:.*]]: tensor) -> !llvm.ptr // CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index // CHECK-DAG: %[[C1:.*]] = arith.constant 1 : index +// CHECK-DAG: %[[U:.*]] = tensor.dim %[[A]], %[[C0]] : tensor // CHECK-DAG: %[[P:.*]] = memref.alloca() : memref<1xi8> // CHECK-DAG: %[[Q:.*]] = memref.alloca() : memref<1xi64> // CHECK-DAG: %[[R:.*]] = memref.alloca() : memref<1xi64> // CHECK-DAG: %[[X:.*]] = memref.cast %[[P]] : memref<1xi8> to memref // CHECK-DAG: %[[Y:.*]] = memref.cast %[[Q]] : memref<1xi64> to memref // CHECK-DAG: %[[Z:.*]] = memref.cast %[[R]] : memref<1xi64> to memref -// CHECK: %[[C:.*]] = call @newSparseTensor(%[[X]], %[[Y]], %[[Z]], %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %{{.}}) +// CHECK: %[[NP:.*]] = llvm.mlir.null : !llvm.ptr +// CHECK: %[[C:.*]] = call @newSparseTensor(%[[X]], %[[Y]], %[[Z]], %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %[[NP]]) // CHECK: %[[M:.*]] = memref.alloca() : memref<1xindex> // CHECK: %[[T:.*]] = memref.cast %[[M]] : memref<1xindex> to memref -// CHECK: %[[U:.*]] = tensor.dim %[[A]], %[[C0]] : tensor // CHECK: scf.for %[[I:.*]] = %[[C0]] to %[[U]] step %[[C1]] { // CHECK: %[[E:.*]] = tensor.extract %[[A]][%[[I]]] : tensor // CHECK: memref.store %[[I]], %[[M]][%[[C0]]] : memref<1xindex> // CHECK: call @addEltI32(%[[C]], %[[E]], %[[T]], %[[Z]]) // CHECK: } -// CHECK: %[[T:.*]] = call @newSparseTensor(%{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %[[C]]) +// CHECK: %[[T:.*]] = call @newSparseTensor(%[[X]], %[[Y]], %[[Z]], %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %[[C]]) // CHECK: return %[[T]] : !llvm.ptr func @sparse_convert_1d(%arg0: tensor) -> tensor { %0 = sparse_tensor.convert %arg0 : tensor to tensor @@ -180,8 +181,14 @@ // CHECK-LABEL: func @sparse_convert_1d_ss( // CHECK-SAME: %[[A:.*]]: !llvm.ptr) -// CHECK: %[[C:.*]] = call @newSparseTensor(%{{.}}, %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %[[A]]) -// CHECK: %[[T:.*]] = call @newSparseTensor(%{{.}}, %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %[[C]]) +// CHECK-DAG: %[[P:.*]] = memref.alloca() : memref<1xi8> +// CHECK-DAG: %[[Q:.*]] = memref.alloca() : memref<1xi64> +// CHECK-DAG: %[[R:.*]] = memref.alloca() : memref<1xi64> +// CHECK-DAG: %[[X:.*]] = memref.cast %[[P]] : memref<1xi8> to memref +// CHECK-DAG: %[[Y:.*]] = memref.cast %[[Q]] : memref<1xi64> to memref +// CHECK-DAG: %[[Z:.*]] = memref.cast %[[R]] : memref<1xi64> to memref +// CHECK: %[[C:.*]] = call @newSparseTensor(%[[X]], %[[Y]], %[[Z]], %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %[[A]]) +// CHECK: %[[T:.*]] = call @newSparseTensor(%[[X]], %[[Y]], %[[Z]], %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %[[C]]) // CHECK: return %[[T]] : !llvm.ptr func @sparse_convert_1d_ss(%arg0: tensor) -> tensor { %0 = sparse_tensor.convert %arg0 : tensor to tensor @@ -198,7 +205,8 @@ // CHECK-DAG: %[[X:.*]] = memref.cast %[[P]] : memref<2xi8> to memref // CHECK-DAG: %[[Y:.*]] = memref.cast %[[Q]] : memref<2xi64> to memref // CHECK-DAG: %[[Z:.*]] = memref.cast %[[R]] : memref<2xi64> to memref -// CHECK: %[[C:.*]] = call @newSparseTensor(%[[X]], %[[Y]], %[[Z]], %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %{{.}}) +// CHECK: %[[NP:.*]] = llvm.mlir.null : !llvm.ptr +// CHECK: %[[C:.*]] = call @newSparseTensor(%[[X]], %[[Y]], %[[Z]], %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %[[NP]]) // CHECK: %[[M:.*]] = memref.alloca() : memref<2xindex> // CHECK: %[[T:.*]] = memref.cast %[[M]] : memref<2xindex> to memref // CHECK: scf.for %[[I:.*]] = %[[C0]] to %{{.*}} step %[[C1]] { @@ -209,7 +217,7 @@ // CHECK: call @addEltF64(%[[C]], %[[E]], %[[T]], %[[Z]]) // CHECK: } // CHECK: } -// CHECK: %[[T:.*]] = call @newSparseTensor(%{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %[[C]]) +// CHECK: %[[T:.*]] = call @newSparseTensor(%[[X]], %[[Y]], %[[Z]], %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %[[C]]) // CHECK: return %[[T]] : !llvm.ptr func @sparse_convert_2d(%arg0: tensor<2x4xf64>) -> tensor<2x4xf64, #SparseMatrix> { %0 = sparse_tensor.convert %arg0 : tensor<2x4xf64> to tensor<2x4xf64, #SparseMatrix> @@ -226,7 +234,8 @@ // CHECK-DAG: %[[X:.*]] = memref.cast %[[P]] : memref<2xi8> to memref // CHECK-DAG: %[[Y:.*]] = memref.cast %[[Q]] : memref<2xi64> to memref // CHECK-DAG: %[[Z:.*]] = memref.cast %[[R]] : memref<2xi64> to memref -// CHECK: %[[C:.*]] = call @newSparseTensor(%[[X]], %[[Y]], %[[Z]], %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %{{.}}) +// CHECK: %[[NP:.*]] = llvm.mlir.null : !llvm.ptr +// CHECK: %[[C:.*]] = call @newSparseTensor(%[[X]], %[[Y]], %[[Z]], %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %[[NP]]) // CHECK: %[[M:.*]] = memref.alloca() : memref<2xindex> // CHECK: %[[N:.*]] = memref.cast %[[M]] : memref<2xindex> to memref // CHECK: scf.for %[[I:.*]] = %[[C0]] to %[[C2]] step %[[C1]] { @@ -235,7 +244,7 @@ // CHECK: %[[V:.*]] = tensor.extract %{{.*}}[%[[I]]] : tensor<2xf32> // CHECK: call @addEltF32(%{{.*}}, %[[V]], %[[N]], %{{.*}}) // CHECK: } -// CHECK: %[[T:.*]] = call @newSparseTensor(%{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %[[C]]) +// CHECK: %[[T:.*]] = call @newSparseTensor(%[[X]], %[[Y]], %[[Z]], %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %[[C]]) // CHECK: return %[[T]] : !llvm.ptr func @sparse_constant() -> tensor<8x7xf32, #SparseMatrix>{ // Initialize a tensor. @@ -250,18 +259,19 @@ // CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index // CHECK-DAG: %[[C1:.*]] = arith.constant 1 : index // CHECK-DAG: %[[C2:.*]] = arith.constant 2 : index +// CHECK-DAG: %[[U1:.*]] = tensor.dim %[[A]], %[[C0]] : tensor +// CHECK-DAG: %[[U2:.*]] = tensor.dim %[[A]], %[[C1]] : tensor +// CHECK-DAG: %[[U3:.*]] = tensor.dim %[[A]], %[[C2]] : tensor // CHECK-DAG: %[[P:.*]] = memref.alloca() : memref<3xi8> // CHECK-DAG: %[[Q:.*]] = memref.alloca() : memref<3xi64> // CHECK-DAG: %[[R:.*]] = memref.alloca() : memref<3xi64> // CHECK-DAG: %[[X:.*]] = memref.cast %[[P]] : memref<3xi8> to memref // CHECK-DAG: %[[Y:.*]] = memref.cast %[[Q]] : memref<3xi64> to memref // CHECK-DAG: %[[Z:.*]] = memref.cast %[[R]] : memref<3xi64> to memref -// CHECK: %[[C:.*]] = call @newSparseTensor(%[[X]], %[[Y]], %[[Z]], %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %{{.}}) +// CHECK: %[[NP:.*]] = llvm.mlir.null : !llvm.ptr +// CHECK: %[[C:.*]] = call @newSparseTensor(%[[X]], %[[Y]], %[[Z]], %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %[[NP]]) // CHECK: %[[M:.*]] = memref.alloca() : memref<3xindex> // CHECK: %[[N:.*]] = memref.cast %[[M]] : memref<3xindex> to memref -// CHECK: %[[U1:.*]] = tensor.dim %[[A]], %[[C0]] : tensor -// CHECK: %[[U2:.*]] = tensor.dim %[[A]], %[[C1]] : tensor -// CHECK: %[[U3:.*]] = tensor.dim %[[A]], %[[C2]] : tensor // CHECK: scf.for %[[I:.*]] = %[[C0]] to %[[U1]] step %[[C1]] { // CHECK: scf.for %[[J:.*]] = %[[C0]] to %[[U2]] step %[[C1]] { // CHECK: scf.for %[[K:.*]] = %[[C0]] to %[[U3]] step %[[C1]] { @@ -273,7 +283,7 @@ // CHECK: } // CHECK: } // CHECK: } -// CHECK: %[[T:.*]] = call @newSparseTensor(%{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %[[C]]) +// CHECK: %[[T:.*]] = call @newSparseTensor(%[[X]], %[[Y]], %[[Z]], %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %[[C]]) // CHECK: return %[[T]] : !llvm.ptr func @sparse_convert_3d(%arg0: tensor) -> tensor { %0 = sparse_tensor.convert %arg0 : tensor to tensor 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 @@ -162,8 +162,8 @@ #CSR = #sparse_tensor.encoding<{dimLevelType = ["dense", "compressed"]}> -func @sparse_convert_mismatch(%arg0: tensor<10x10xf32>) -> tensor<10x?xf32, #CSR> { +func @sparse_convert_mismatch(%arg0: tensor<10x?xf32>) -> tensor<10x10xf32, #CSR> { // expected-error@+1 {{unexpected conversion mismatch in dimension 1}} - %0 = sparse_tensor.convert %arg0 : tensor<10x10xf32> to tensor<10x?xf32, #CSR> - return %0 : tensor<10x?xf32, #CSR> + %0 = sparse_tensor.convert %arg0 : tensor<10x?xf32> to tensor<10x10xf32, #CSR> + return %0 : tensor<10x10xf32, #CSR> } diff --git a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_convert.mlir b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_convert.mlir new file mode 100644 --- /dev/null +++ b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_convert.mlir @@ -0,0 +1,91 @@ +// RUN: mlir-opt %s \ +// RUN: --sparsification --sparse-tensor-conversion \ +// RUN: --linalg-bufferize --convert-linalg-to-loops \ +// RUN: --convert-vector-to-scf --convert-scf-to-std \ +// RUN: --func-bufferize --tensor-constant-bufferize --tensor-bufferize \ +// RUN: --std-bufferize --finalizing-bufferize --lower-affine \ +// RUN: --convert-vector-to-llvm --convert-memref-to-llvm --convert-math-to-llvm \ +// RUN: --convert-std-to-llvm --reconcile-unrealized-casts | \ +// RUN: mlir-cpu-runner \ +// RUN: -e entry -entry-point-result=void \ +// RUN: -shared-libs=%mlir_integration_test_dir/libmlir_c_runner_utils%shlibext | \ +// RUN: FileCheck %s + +#DCSR = #sparse_tensor.encoding<{ + dimLevelType = [ "compressed", "compressed" ] +}> + +#DCSC = #sparse_tensor.encoding<{ + dimLevelType = [ "compressed", "compressed" ], + dimOrdering = affine_map<(i,j) -> (j,i)> +}> + +// +// Integration test that tests conversions between sparse tensors, +// where the dynamic sizes of the shape of the enveloping tensor +// may change (the actual underlying sizes obviously never change). +// +module { + + // + // Helper method to print values array. The transfer actually + // reads more than required to verify size of buffer as well. + // + func @dump(%arg0: memref) { + %c = arith.constant 0 : index + %d = arith.constant -1.0 : f64 + %0 = vector.transfer_read %arg0[%c], %d: memref, vector<8xf64> + vector.print %0 : vector<8xf64> + return + } + + func @entry() { + %t1 = arith.constant sparse< + [ [0,0], [0,1], [0,63], [1,0], [1,1], [31,0], [31,63] ], + [ 1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0 ]> : tensor<32x64xf64> + %t2 = tensor.cast %t1 : tensor<32x64xf64> to tensor + + // Four dense to sparse conversions. + %1 = sparse_tensor.convert %t1 : tensor<32x64xf64> to tensor + %2 = sparse_tensor.convert %t1 : tensor<32x64xf64> to tensor + %3 = sparse_tensor.convert %t2 : tensor to tensor + %4 = sparse_tensor.convert %t2 : tensor to tensor + + // Two cross conversions. + %5 = sparse_tensor.convert %3 : tensor to tensor + %6 = sparse_tensor.convert %4 : tensor to tensor + + // + // All proper row-/column-wise? + // + // CHECK: ( 1, 2, 3, 4, 5, 6, 7, -1 ) + // CHECK: ( 1, 4, 6, 2, 5, 3, 7, -1 ) + // CHECK: ( 1, 2, 3, 4, 5, 6, 7, -1 ) + // CHECK: ( 1, 4, 6, 2, 5, 3, 7, -1 ) + // CHECK: ( 1, 4, 6, 2, 5, 3, 7, -1 ) + // CHECK: ( 1, 2, 3, 4, 5, 6, 7, -1 ) + // + %m1 = sparse_tensor.values %1 : tensor to memref + %m2 = sparse_tensor.values %2 : tensor to memref + %m3 = sparse_tensor.values %3 : tensor to memref + %m4 = sparse_tensor.values %4 : tensor to memref + %m5 = sparse_tensor.values %5 : tensor to memref + %m6 = sparse_tensor.values %6 : tensor to memref + call @dump(%m1) : (memref) -> () + call @dump(%m2) : (memref) -> () + call @dump(%m3) : (memref) -> () + call @dump(%m4) : (memref) -> () + call @dump(%m5) : (memref) -> () + call @dump(%m6) : (memref) -> () + + // Release the resources. + sparse_tensor.release %1 : tensor + sparse_tensor.release %2 : tensor + sparse_tensor.release %3 : tensor + sparse_tensor.release %4 : tensor + sparse_tensor.release %5 : tensor + sparse_tensor.release %6 : tensor + + return + } +}