[mlir][sparse] Refactor the convert operator conversion to support codegen for the operator.

[bixia]

Summary

Outline the code that generates the loop structure to iterate over a dense
tensor or a sparse constant to genDenseTensorOrSparseConstantIterLoop.

Move a few routines to CodegenUtils for sharing.

[Peiming]

mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorConversion.cpp

816

Is the variable used in this function?

[bixia]

Rebase.

[bixia]

Remove an unused variable.

[bixia]

mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorConversion.cpp

816

Good catch. Fixed.

[bixia]

Rebase.

[wrengr]

mlir/lib/Dialect/SparseTensor/Transforms/CodegenUtils.cpp

80

"an"

mlir/lib/Dialect/SparseTensor/Transforms/CodegenUtils.h

269

Should have a space after the period

[bixia]

Fix comment.

mlir/lib/Dialect/SparseTensor/Transforms/CodegenUtils.h

 /// type, but returns it as type `memref<? x $tp>` (rather than as type
 /// `memref<$sz x $tp>`).
 Value genAlloca(OpBuilder &builder, Location loc, unsigned sz, Type tp);
 
 /// Generates an uninitialized temporary buffer with room for one value
 /// of the given type, and returns the `memref<$tp>`.
 Value genAllocaScalar(OpBuilder &builder, Location loc, Type tp);
 
+/// Generates code to allocate a buffer of the given type, and zero
+/// initialize it.  If the buffer type has any dynamic sizes, then the
+/// `sizes` parameter should be as filled by sizesFromPtr(); that way
+/// we can reuse the genDimSizeCall() results generated by sizesFromPtr().
+Value allocDenseTensor(OpBuilder &builder, Location loc,
+                       RankedTensorType tensorTp, ValueRange sizes);
+
+/// Generates the code to read the value from tensor[ivs]. The generated code
+/// looks like the following and the insertion point after this routine is
+/// inside the if-then branch behind the assignment to ind.
+///    if (tensor[ivs] != 0)
+///      insert_point
+Value genValueForDense(OpBuilder &builder, Location loc, Value tensor,
+                       ValueRange ivs);
+
+/// Generates the loop structure to iterate over a dense tensor or a sparse
+/// tensor constant to support the lowering of dense-to-sparse convert operator.
+//
+// The loop to iterate a dense tensor:
+//   for i1 in dim1
+//    ..
+//     for ik in dimk
+//       val = a[i1,..,ik]
+//       if val != 0
+//         loop-body
+//
+// The loop to iterate a sparse tensor constant:
+//   for i in range(NNZ)
+//     val = values[i]
+//     [i1,..,ik] = indices[i]
+//     loop-body
+void genDenseTensorOrSparseConstantIterLoop(
+    OpBuilder &builder, Location loc, Value src, unsigned rank,
+    function_ref<void(OpBuilder &, Location, Value, ValueRange)> bodyBuilder);
+
+/// Populates given sizes array from dense tensor or sparse tensor constant.
+void sizesFromSrc(OpBuilder &builder, SmallVector<Value, 4> &sizes,
+                  Location loc, Value src);
+
 //===----------------------------------------------------------------------===//
 // Inlined constant generators.
 //
 // All these functions are just wrappers to improve code legibility;
 // therefore, we mark them as `inline` to avoid introducing any additional
 // overhead due to the legibility.
 //
 // TODO: Ideally these should move upstream, so that we don't

mlir/lib/Dialect/SparseTensor/Transforms/CodegenUtils.cpp

 //===- CodegenUtils.cpp - Utilities for generating MLIR -------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "CodegenUtils.h"
 
 #include "mlir/Dialect/Affine/IR/AffineOps.h"
 #include "mlir/Dialect/Bufferization/IR/Bufferization.h"
 #include "mlir/Dialect/Linalg/IR/Linalg.h"
 #include "mlir/Dialect/Linalg/Utils/Utils.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
+#include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "mlir/IR/Matchers.h"
 #include "mlir/IR/Types.h"
 #include "mlir/IR/Value.h"
 
 using namespace mlir;
 using namespace mlir::sparse_tensor;
 
 /// Generates a pointer/index load from the sparse storage scheme. Narrower
if (!load.getType().isa<IndexType>()) {
     if (load.getType().getIntOrFloatBitWidth() < 64)
       load = builder.create<arith::ExtUIOp>(loc, builder.getI64Type(), load);
     load =
         builder.create<arith::IndexCastOp>(loc, builder.getIndexType(), load);
   }
   return load;
 }
 
+/// If the tensor is a sparse constant, generates and returns the pair of
+/// the constants for the indices and the values.
+static Optional<std::pair<Value, Value>>
+genSplitSparseConstant(OpBuilder &builder, Location loc, Value tensor) {
+  if (auto constOp = tensor.getDefiningOp<arith::ConstantOp>()) {
+    if (auto attr = constOp.getValue().dyn_cast<SparseElementsAttr>()) {
+      DenseElementsAttr indicesAttr = attr.getIndices();
+      Value indices = builder.create<arith::ConstantOp>(loc, indicesAttr);
+      DenseElementsAttr valuesAttr = attr.getValues();
+      Value values = builder.create<arith::ConstantOp>(loc, valuesAttr);
+      return std::make_pair(indices, values);
+    }
+  }
+  return {};
+}
+
+/// Generates the code to copy the index at indices[ivs] to ind, and return
+/// the value at value[ivs].
+static Value genIndexAndValueForSparse(OpBuilder &builder, Location loc,
+                                       Value indices, Value values,
+                                       SmallVectorImpl<Value> &indicesArray,
+                                       ValueRange ivs, unsigned rank) {
+  for (unsigned i = 0; i < rank; i++) {
+    Value idx = constantIndex(builder, loc, i);
+    Value val = builder.create<tensor::ExtractOp>(loc, indices,
+                                                  ValueRange{ivs[0], idx});
+    val = builder.create<arith::IndexCastOp>(loc, builder.getIndexType(), val);
+    // builder.create<memref::StoreOp>(loc, val, ind, idx);
+    indicesArray.push_back(val);
+  }
+  return builder.create<tensor::ExtractOp>(loc, values, ivs[0]);
+}
+
+/// Generates the code to read the value from tensor[ivs], and conditionally
+/// stores the indices ivs to the memory in ind. The generated code looks like
+/// the following and the insertion point after this routine is inside the
+/// if-then branch behind the assignment to ind. This is to ensure that the
+/// code that uses the ind, such as an addEltX call generated after, is inside
+/// the if-then branch.
+///    if (tensor[ivs] != 0)
+///      ind = ivs
+static Value genIndexAndValueForDense(OpBuilder &builder, Location loc,
+                                      Value tensor,
+                                      SmallVectorImpl<Value> &indicesArray,
+                                      ValueRange ivs) {
+  Value val = genValueForDense(builder, loc, tensor, ivs);
+  indicesArray.append(ivs.begin(), ivs.end());
+  return val;
+}
+
 //===----------------------------------------------------------------------===//
 // Sparse tensor loop emitter class implementations
 //===----------------------------------------------------------------------===//
 
 SparseTensorLoopEmitter::SparseTensorLoopEmitter(ValueRange tensors,
                                                  bool isLastOutput)
     : tensors(tensors.begin(), tensors.end()), dims(tensors.size()),
       pidxs(tensors.size()), coord(tensors.size()), highs(tensors.size()),
Value mlir::sparse_tensor::genAlloca(OpBuilder &builder, Location loc, Value sz,
   auto memTp = MemRefType::get({ShapedType::kDynamicSize}, tp);
   return builder.create<memref::AllocaOp>(loc, memTp, ValueRange{sz});
 }
 
 Value mlir::sparse_tensor::genAllocaScalar(OpBuilder &builder, Location loc,
                                            Type tp) {
   return builder.create<memref::AllocaOp>(loc, MemRefType::get({}, tp));
 }
+
+Value mlir::sparse_tensor::allocDenseTensor(OpBuilder &builder, Location loc,
+                                            RankedTensorType tensorTp,
+                                            ValueRange sizes) {
+  Type elemTp = tensorTp.getElementType();
+  auto shape = tensorTp.getShape();
+  auto memTp = MemRefType::get(shape, elemTp);
+  SmallVector<Value> dynamicSizes;
+  for (unsigned i = 0, rank = tensorTp.getRank(); i < rank; i++) {
+    if (shape[i] == ShapedType::kDynamicSize)
+      dynamicSizes.push_back(sizes[i]);
+  }
+  Value mem = builder.create<memref::AllocOp>(loc, memTp, dynamicSizes);
+  Value zero = constantZero(builder, loc, elemTp);
+  builder.create<linalg::FillOp>(loc, ValueRange{zero}, ValueRange{mem});
+  return mem;
+}
+
+Value mlir::sparse_tensor::genValueForDense(OpBuilder &builder, Location loc,
+                                            Value tensor, ValueRange ivs) {
+  Value val = builder.create<tensor::ExtractOp>(loc, tensor, ivs);
+  Value cond = genIsNonzero(builder, loc, val);
+  scf::IfOp ifOp = builder.create<scf::IfOp>(loc, cond, /*else*/ false);
+  builder.setInsertionPointToStart(&ifOp.getThenRegion().front());
+  return val;
+}
+
+void mlir::sparse_tensor::genDenseTensorOrSparseConstantIterLoop(
+    OpBuilder &builder, Location loc, Value src, unsigned rank,
+    function_ref<void(OpBuilder &, Location, Value, ValueRange)> bodyBuilder) {
+  SmallVector<Value, 4> indicesArray;
+  SmallVector<Value> lo;
+  SmallVector<Value> hi;
+  SmallVector<Value> st;
+  Value zero = constantIndex(builder, loc, 0);
+  Value one = constantIndex(builder, loc, 1);
+  auto indicesValues = genSplitSparseConstant(builder, loc, src);
+  bool isCOOConstant = indicesValues.has_value();
+  Value indices;
+  Value values;
+  if (isCOOConstant) {
+    indices = indicesValues->first;
+    values = indicesValues->second;
+    lo.push_back(zero);
+    hi.push_back(linalg::createOrFoldDimOp(builder, loc, values, 0));
+    st.push_back(one);
+  } else {
+    for (unsigned i = 0; i < rank; i++) {
+      lo.push_back(zero);
+      hi.push_back(linalg::createOrFoldDimOp(builder, loc, src, i));
+      st.push_back(one);
+    }
+  }
+
+  scf::buildLoopNest(
+      builder, loc, lo, hi, st, {},
+      [&](OpBuilder &builder, Location loc, ValueRange ivs,
+          ValueRange args) -> scf::ValueVector {
+        Value val;
+        if (isCOOConstant)
+          val = genIndexAndValueForSparse(builder, loc, indices, values,
+                                          indicesArray, ivs, rank);
+        else
+          val = genIndexAndValueForDense(builder, loc, src, indicesArray, ivs);
+        bodyBuilder(builder, loc, val, indicesArray);
+        return {};
+      });
+}
+
+void mlir::sparse_tensor::sizesFromSrc(OpBuilder &builder,
+                                       SmallVector<Value, 4> &sizes,
+                                       Location loc, Value src) {
+  unsigned rank = src.getType().cast<ShapedType>().getRank();
+  for (unsigned i = 0; i < rank; i++)
+    sizes.push_back(linalg::createOrFoldDimOp(builder, loc, src, i));
+}

mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorConversion.cpp

static void sizesFromType(OpBuilder &builder, SmallVector<Value, 4> &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(builder, loc, s));
   }
 }
 
-/// Populates given sizes array from source.
-static void sizesFromSrc(OpBuilder &builder, SmallVector<Value, 4> &sizes,
-                         Location loc, Value src) {
-  unsigned rank = src.getType().cast<ShapedType>().getRank();
-  for (unsigned i = 0; i < rank; i++)
-    sizes.push_back(linalg::createOrFoldDimOp(builder, loc, src, i));
-}
-
 /// Populates the given sizes array for concatenation from type (for static
 /// sizes) and from an already-converted opaque pointer source (for dynamic
 /// sizes).
 static void concatSizesFromInputs(OpBuilder &builder,
                                   SmallVector<Value, 4> &sizes, Location loc,
                                   ShapedType dstTp, ValueRange srcs,
                                   unsigned dim) {
   auto dstShape = dstTp.getShape();
static void newParams(OpBuilder &builder, SmallVector<Value, 8> &params,
   // User action.
   params.push_back(constantAction(builder, loc, action));
   // Payload pointer.
   if (!ptr)
     ptr = builder.create<LLVM::NullOp>(loc, getOpaquePointerType(builder));
   params.push_back(ptr);
 }
 
-/// Generates the code to read the value from tensor[ivs].The generated code
-/// looks like the following and the insertion point after this routine is
-/// inside the if-then branch behind the assignment to ind.
-///    if (tensor[ivs] != 0)
-///      insert_point
-static Value genValueForDense(OpBuilder &builder, Location loc, Value tensor,
-                              ValueRange ivs) {
-  Value val = builder.create<tensor::ExtractOp>(loc, tensor, ivs);
-  Value cond = genIsNonzero(builder, loc, val);
-  scf::IfOp ifOp = builder.create<scf::IfOp>(loc, cond, /*else*/ false);
-  builder.setInsertionPointToStart(&ifOp.getThenRegion().front());
-  return val;
-}
-
-/// Generates the code to read the value from tensor[ivs], and conditionally
-/// stores the indices ivs to the memory in ind. The generated code looks like
-/// the following and the insertion point after this routine is inside the
-/// if-then branch behind the assignment to ind. This is to ensure that the
-/// addEltX call generated after is inside the if-then branch.
-///    if (tensor[ivs] != 0)
-///      ind = ivs
-static Value genIndexAndValueForDense(OpBuilder &builder, Location loc,
-                                      Value tensor, Value ind, ValueRange ivs) {
-  Value val = genValueForDense(builder, loc, tensor, ivs);
-  unsigned i = 0;
-  for (auto iv : ivs) {
-    Value idx = constantIndex(builder, loc, i++);
-    builder.create<memref::StoreOp>(loc, iv, ind, idx);
-  }
-  return val;
-}
-
 /// Generates a call to release/delete a `SparseTensorCOO`.
 static void genDelCOOCall(OpBuilder &builder, Location loc, Type elemTp,
                           Value coo) {
   SmallString<21> name{"delSparseTensorCOO", primaryTypeFunctionSuffix(elemTp)};
   createFuncCall(builder, loc, name, {}, coo, EmitCInterface::Off);
 }
 
 /// Generates a call to release/delete a `SparseTensorIterator`.
static Value genGetNextCall(OpBuilder &builder, Location loc, Value iter,
   Type elemTp = elemPtr.getType().cast<ShapedType>().getElementType();
   SmallString<10> name{"getNext", primaryTypeFunctionSuffix(elemTp)};
   SmallVector<Value, 3> params{iter, ind, elemPtr};
   Type i1 = builder.getI1Type();
   return createFuncCall(builder, loc, name, i1, params, EmitCInterface::On)
       .getResult(0);
 }
 
-/// If the tensor is a sparse constant, generates and returns the pair of
-/// the constants for the indices and the values.
-static Optional<std::pair<Value, Value>>
-genSplitSparseConstant(OpBuilder &builder, Location loc, Value tensor) {
-  if (auto constOp = tensor.getDefiningOp<arith::ConstantOp>()) {
-    if (auto attr = constOp.getValue().dyn_cast<SparseElementsAttr>()) {
-      DenseElementsAttr indicesAttr = attr.getIndices();
-      Value indices = builder.create<arith::ConstantOp>(loc, indicesAttr);
-      DenseElementsAttr valuesAttr = attr.getValues();
-      Value values = builder.create<arith::ConstantOp>(loc, valuesAttr);
-      return std::make_pair(indices, values);
-    }
-  }
-  return {};
-}
-
-/// Generates the code to copy the index at indices[ivs] to ind, and return
-/// the value at value[ivs].
-static Value genIndexAndValueForSparse(OpBuilder &builder, Location loc,
-                                       Value indices, Value values, Value ind,
-                                       ValueRange ivs, unsigned rank) {
-  for (unsigned i = 0; i < rank; i++) {
-    Value idx = constantIndex(builder, loc, i);
-    Value val = builder.create<tensor::ExtractOp>(loc, indices,
-                                                  ValueRange{ivs[0], idx});
-    val = builder.create<arith::IndexCastOp>(loc, builder.getIndexType(), val);
-    builder.create<memref::StoreOp>(loc, val, ind, idx);
-  }
-  return builder.create<tensor::ExtractOp>(loc, values, ivs[0]);
-}
-
-/// Generates code to allocate a buffer of the given type, and zero
-/// initialize it.  If the buffer type has any dynamic sizes, then the
-/// `sizes` parameter should be as filled by sizesFromPtr(); that way
-/// we can reuse the genDimSizeCall() results generated by sizesFromPtr().
-static Value allocDenseTensor(OpBuilder &builder, Location loc,
-                              RankedTensorType tensorTp, ValueRange sizes) {
-  Type elemTp = tensorTp.getElementType();
-  auto shape = tensorTp.getShape();
-  auto memTp = MemRefType::get(shape, elemTp);
-  SmallVector<Value> dynamicSizes;
-  for (unsigned i = 0, rank = tensorTp.getRank(); i < rank; i++) {
-    if (shape[i] == ShapedType::kDynamicSize)
-      dynamicSizes.push_back(sizes[i]);
-  }
-  Value mem = builder.create<memref::AllocOp>(loc, memTp, dynamicSizes);
-  Value zero = constantZero(builder, loc, elemTp);
-  builder.create<linalg::FillOp>(loc, ValueRange{zero}, ValueRange{mem});
-  return mem;
-}
-
 /// Generates code to deallocate a dense buffer.
 static void deallocDenseTensor(OpBuilder &builder, Location loc, Value buffer) {
   builder.create<memref::DeallocOp>(loc, buffer);
 }
 
 /// Converts a pointer to COO (from calls to iter->next()) into a vector of
 /// indices, apply (optional) `offset` on `offsetDim`.
 static SmallVector<Value, 4> loadIndices(OpBuilder &builder, Location loc,
matchAndRewrite(ConvertOp op, OpAdaptor adaptor,
     ShapedType stp = resType.cast<ShapedType>();
     unsigned rank = stp.getRank();
     SmallVector<Value, 4> sizes;
     SmallVector<Value, 8> params;
     sizesFromSrc(rewriter, sizes, loc, src);
     newParams(rewriter, params, loc, stp, encDst, Action::kEmptyCOO, sizes);
     Value coo = genNewCall(rewriter, loc, params);
     Value ind = genAlloca(rewriter, loc, rank, rewriter.getIndexType());
     Value perm = params[2];
-    SmallVector<Value> lo;
-    SmallVector<Value> hi;
-    SmallVector<Value> st;
-    Value zero = constantIndex(rewriter, loc, 0);
-    Value one = constantIndex(rewriter, loc, 1);
-    auto indicesValues = genSplitSparseConstant(rewriter, loc, src);
-    bool isCOOConstant = indicesValues.has_value();
-    Value indices;
-    Value values;
-    if (isCOOConstant) {
-      indices = indicesValues->first;
-      values = indicesValues->second;
-      lo.push_back(zero);
-      hi.push_back(linalg::createOrFoldDimOp(rewriter, loc, values, 0));
-      st.push_back(one);
-    } else {
-      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 = stp.getElementType();
     Value elemPtr = genAllocaScalar(rewriter, loc, eltType);
-    scf::buildLoopNest(
-        rewriter, op.getLoc(), lo, hi, st, {},
-        [&](OpBuilder &builder, Location loc, ValueRange ivs,
-            ValueRange args) -> scf::ValueVector {
-          Value val;
-          if (isCOOConstant)
-            val = genIndexAndValueForSparse(rewriter, loc, indices, values, ind,
-                                            ivs, rank);
-          else
-            val = genIndexAndValueForDense(rewriter, loc, src, ind, ivs);
+    genDenseTensorOrSparseConstantIterLoop(
+        rewriter, loc, src, rank,
+        [&](OpBuilder &builder, Location loc, Value val, ValueRange indices) {
+          for (unsigned i = 0; i < rank; i++) {
+            Value idx = constantIndex(builder, loc, i);
+            builder.create<memref::StoreOp>(loc, indices[i], ind, idx);
+          }
           builder.create<memref::StoreOp>(loc, val, elemPtr);
-          genAddEltCall(rewriter, loc, eltType, coo, elemPtr, ind, perm);
-          return {};
+          genAddEltCall(builder, loc, eltType, coo, elemPtr, ind, perm);
         });
     // Final call to construct sparse tensor storage.
     params[6] = constantAction(rewriter, loc, Action::kFromCOO);
     params[7] = coo;
     Value dst = genNewCall(rewriter, loc, params);
     genDelCOOCall(rewriter, loc, eltType, coo);
     rewriter.replaceOp(op, dst);
     return success();

mlir/test/Dialect/SparseTensor/convert_dense2sparse.mlir

 //   CHECK-DAG: %[[Y:.*]] = memref.cast %[[Q]] : memref<2xindex> to memref<?xindex>
 //   CHECK-DAG: %[[Z:.*]] = memref.cast %[[R]] : memref<2xindex> to memref<?xindex>
 //       CHECK: %[[NP:.*]] = llvm.mlir.null : !llvm.ptr<i8>
 //       CHECK: %[[C:.*]] = call @newSparseTensor(%[[X]], %[[Y]], %[[Z]], %{{.*}}, %{{.*}}, %{{.*}}, %[[EmptyCOO]], %[[NP]])
 //       CHECK: %[[M:.*]] = memref.alloca() : memref<2xindex>
 //       CHECK: %[[N:.*]] = memref.cast %[[M]] : memref<2xindex> to memref<?xindex>
 //       CHECK: %[[BUF:.*]] = memref.alloca() : memref<f32>
 //       CHECK: scf.for %[[I:.*]] = %[[C0]] to %[[C2]] step %[[C1]] {
-//       CHECK:   memref.store %{{.*}}, %[[M]][%[[C0]]] : memref<2xindex>
-//       CHECK:   memref.store %{{.*}}, %[[M]][%[[C1]]] : memref<2xindex>
-//       CHECK:   %[[V:.*]] = tensor.extract %{{.*}}[%[[I]]] : tensor<2xf32>
+//       CHECK-DAG:   memref.store %{{.*}}, %[[M]][%[[C0]]] : memref<2xindex>
+//       CHECK-DAG:   memref.store %{{.*}}, %[[M]][%[[C1]]] : memref<2xindex>
+//       CHECK-DAG:   %[[V:.*]] = tensor.extract %{{.*}}[%[[I]]] : tensor<2xf32>
 //       CHECK:   memref.store %[[V]], %[[BUF]][] : memref<f32>
 //       CHECK:   call @addEltF32(%{{.*}}, %[[BUF]], %[[N]], %{{.*}})
 //       CHECK: }
 //       CHECK: %[[T:.*]] = call @newSparseTensor(%[[X]], %[[Y]], %[[Z]], %{{.*}}, %{{.*}}, %{{.*}}, %[[FromCOO]], %[[C]])
 //       CHECK: call @delSparseTensorCOOF32(%[[C]])
 //       CHECK: return %[[T]] : !llvm.ptr<i8>
 func.func @sparse_constant() -> tensor<8x7xf32, #CSR>{
   // Initialize a tensor.