[mlir][sparse] Code cleanup for SparseTensorConversion

[aartbik]

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

145

Add a /// comment

174

same

181

same

216

/*emitCInterface=*/
was dropped

320–321

what does this mean?

392

/*emitCInterface=*/t

798

/*emitCInterface=*/t

826

/*emitCInterface=*/t

856

/*emitCInterface=*/t

905

/*emitCInterface=*/t

[wrengr]

Addressing comments

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

320–321

I mean that if the elemTp is not one we support then constantPrimaryTypeEncoding will call llvm_unreachable. This isn't a change from previous behavior, just a comment about it. Really it's just a note-to-self I forgot to remove re the same concern discussed at D115012 re calling llvm_unreachable vs returning failure().

[wrengr]

one more cleanup (inlined single-use Location variable)

[aartbik]

Looking good!

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

case SparseTensorEncodingAttr::DimLevelType::Compressed:
     break;
   case SparseTensorEncodingAttr::DimLevelType::Singleton:
     dlt2 = DimLevelType::kSingleton;
     break;
   }
   return constantI8(rewriter, loc, static_cast<uint8_t>(dlt2));
 }
 
+/// Returns the equivalent of `void*` for opaque arguments to the
+/// execution engine.
+static Type getOpaquePointerType(PatternRewriter &rewriter) {
+  return LLVM::LLVMPointerType::get(rewriter.getI8Type());
+}
+
 /// Returns a function reference (first hit also inserts into module). Sets
 /// the "_emit_c_interface" on the function declaration when requested,
 /// so that LLVM lowering generates a wrapper function that takes care
 /// of ABI complications with passing in and returning MemRefs to C functions.
 static FlatSymbolRefAttr getFunc(Operation *op, StringRef name,
                                  TypeRange resultType, ValueRange operands,
-                                 bool emitCInterface = false) {
+                                 bool emitCInterface) {
   MLIRContext *context = op->getContext();
   auto module = op->getParentOfType<ModuleOp>();
   auto result = SymbolRefAttr::get(context, name);
   auto func = module.lookupSymbol<FuncOp>(result.getAttr());
   if (!func) {
     OpBuilder moduleBuilder(module.getBodyRegion());
     func = moduleBuilder.create<FuncOp>(
         op->getLoc(), name,
         FunctionType::get(context, operands.getTypes(), resultType));
     func.setPrivate();
     if (emitCInterface)
       func->setAttr("llvm.emit_c_interface", UnitAttr::get(context));
   }
   return result;
 }
 
+/// Creates a `CallOp` to the function reference returned by `getFunc()`.
+static CallOp createFuncCall(OpBuilder &builder, Operation *op, StringRef name,
+                             TypeRange resultType, ValueRange operands,
+                             bool emitCInterface = false) {
+  auto fn = getFunc(op, name, resultType, operands, emitCInterface);
+  return builder.create<CallOp>(op->getLoc(), resultType, fn, operands);
+}
+
+/// Replaces the `op` with  a `CallOp` to the function reference returned
+/// by `getFunc()`.
+static CallOp replaceOpWithFuncCall(PatternRewriter &rewriter, Operation *op,
+                                    StringRef name, TypeRange resultType,
+                                    ValueRange operands,
+                                    bool emitCInterface = false) {
+  auto fn = getFunc(op, name, resultType, operands, emitCInterface);
+  return rewriter.replaceOpWithNewOp<CallOp>(op, resultType, fn, operands);
+}
+
 /// 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<Value, 2> params;
-  params.push_back(src);
-  params.push_back(constantIndex(rewriter, loc, idx));
+  SmallVector<Value, 2> params{src, constantIndex(rewriter, loc, idx)};
   Type iTp = rewriter.getIndexType();
-  auto fn = getFunc(op, name, iTp, params);
-  return rewriter.create<CallOp>(loc, iTp, fn, params).getResult(0);
+  return createFuncCall(rewriter, op, name, iTp, 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<Value> 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<CallOp>(loc, pTp, fn, params);
+  Type pTp = getOpaquePointerType(rewriter);
+  auto call = createFuncCall(rewriter, op, name, pTp, params,
+                             /*emitCInterface=*/true);
   return call.getResult(0);
 }
 
 /// Populates given sizes array from type.
 static void sizesFromType(ConversionPatternRewriter &rewriter,
                           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(rewriter, loc, s));
   }
 }
 
 /// Populates given sizes array from source.
 static void sizesFromSrc(ConversionPatternRewriter &rewriter,
                          SmallVector<Value, 4> &sizes, Location loc,
                          Value src) {
-  ShapedType stp = src.getType().cast<ShapedType>();
-  for (unsigned i = 0, rank = stp.getRank(); i < rank; i++)
+  unsigned rank = src.getType().cast<ShapedType>().getRank();
+  for (unsigned i = 0; 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<Value, 4> &sizes, Operation *op,
                          SparseTensorEncodingAttr &enc, ShapedType stp,
                          Value src) {
+  Location loc = op->getLoc();
   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]));
+      sizes.push_back(constantIndex(rewriter, loc, shape[i]));
 }
 
 /// Generates an uninitialized temporary buffer of the given size and
 /// type, but returns it as type `memref<? x $tp>` (rather than as type
 /// `memref<$sz x $tp>`).
 static Value genAlloca(ConversionPatternRewriter &rewriter, Location loc,
                        unsigned sz, Type tp) {
   auto memTp = MemRefType::get({ShapedType::kDynamicSize}, tp);
if (AffineMap p = enc.getDimOrdering()) {
     for (unsigned i = 0; i < sz; i++)
       rev[p.getDimPosition(i)] = constantIndex(rewriter, loc, i);
   } else {
     for (unsigned i = 0; i < sz; i++)
       rev[i] = constantIndex(rewriter, loc, i);
   }
   params.push_back(genBuffer(rewriter, loc, rev));
   // Secondary and primary types encoding.
-  ShapedType resType = op->getResult(0).getType().cast<ShapedType>();
+  Type elemTp = op->getResult(0).getType().cast<ShapedType>().getElementType();
   params.push_back(constantPointerTypeEncoding(rewriter, loc, enc));
   params.push_back(constantIndexTypeEncoding(rewriter, loc, enc));
-  params.push_back(
-      constantPrimaryTypeEncoding(rewriter, loc, resType.getElementType()));
-  // User action and pointer.
-  Type pTp = LLVM::LLVMPointerType::get(rewriter.getI8Type());
-  if (!ptr)
-    ptr = rewriter.create<LLVM::NullOp>(loc, pTp);
+  params.push_back(constantPrimaryTypeEncoding(rewriter, loc, elemTp));
+  // User action.
   params.push_back(constantAction(rewriter, loc, action));
+  // Payload pointer.
+  if (!ptr)
+    ptr = rewriter.create<LLVM::NullOp>(loc, getOpaquePointerType(rewriter));
   params.push_back(ptr);
 }
 
 /// Generates the comparison `v != 0` where `v` is of numeric type `t`.
 /// For floating types, we use the "unordered" comparator (i.e., returns
 /// true if `v` is NaN).
 static Value genIsNonzero(ConversionPatternRewriter &rewriter, Location loc,
                           Value v) {

 /// Generates a call that adds one element to a coordinate scheme.
 /// In particular, this generates code like the following:
 ///   val = a[i1,..,ik];
 ///   if val != 0
 ///     t->add(val, [i1,..,ik], [p1,..,pk]);
 static void genAddEltCall(ConversionPatternRewriter &rewriter, Operation *op,
                           Type eltType, Value ptr, Value val, Value ind,
                           Value perm) {
-  Location loc = op->getLoc();
   StringRef name;
   if (eltType.isF64())
     name = "addEltF64";
   else if (eltType.isF32())
     name = "addEltF32";
   else if (eltType.isInteger(64))
     name = "addEltI64";
   else if (eltType.isInteger(32))
     name = "addEltI32";
   else if (eltType.isInteger(16))
     name = "addEltI16";
   else if (eltType.isInteger(8))
     name = "addEltI8";
   else
     llvm_unreachable("Unknown element type");
-  SmallVector<Value, 8> params;
-  params.push_back(ptr);
-  params.push_back(val);
-  params.push_back(ind);
-  params.push_back(perm);
-  Type pTp = LLVM::LLVMPointerType::get(rewriter.getI8Type());
-  auto fn = getFunc(op, name, pTp, params, /*emitCInterface=*/true);
-  rewriter.create<CallOp>(loc, pTp, fn, params);
+  SmallVector<Value, 4> params{ptr, val, ind, perm};
+  Type pTp = getOpaquePointerType(rewriter);
+  createFuncCall(rewriter, op, name, pTp, params, /*emitCInterface=*/true);
 }
 
 /// Generates a call to `iter->getNext()`.  If there is a next element,
 /// then it is copied into the out-parameters `ind` and `elemPtr`,
 /// and the return value is true.  If there isn't a next element, then
 /// the memory for `iter` is freed and the return value is false.
 static Value genGetNextCall(ConversionPatternRewriter &rewriter, Operation *op,
                             Value iter, Value ind, Value elemPtr) {
-  Location loc = op->getLoc();
   Type elemTp = elemPtr.getType().cast<ShapedType>().getElementType();
   StringRef name;
   if (elemTp.isF64())
     name = "getNextF64";
   else if (elemTp.isF32())
     name = "getNextF32";
   else if (elemTp.isInteger(64))
     name = "getNextI64";
   else if (elemTp.isInteger(32))
     name = "getNextI32";
   else if (elemTp.isInteger(16))
     name = "getNextI16";
   else if (elemTp.isInteger(8))
     name = "getNextI8";
   else
     llvm_unreachable("Unknown element type");
-  SmallVector<Value, 3> params;
-  params.push_back(iter);
-  params.push_back(ind);
-  params.push_back(elemPtr);
+  SmallVector<Value, 3> params{iter, ind, elemPtr};
   Type i1 = rewriter.getI1Type();
-  auto fn = getFunc(op, name, i1, params, /*emitCInterface=*/true);
-  auto call = rewriter.create<CallOp>(loc, i1, fn, params);
+  auto call = createFuncCall(rewriter, op, name, i1, params,
+                             /*emitCInterface=*/true);
   return call.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(ConversionPatternRewriter &rewriter, Location loc,
                        Value tensor) {
static Value allocDenseTensor(ConversionPatternRewriter &rewriter, Location loc,
   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 = rewriter.create<memref::AllocOp>(loc, memTp, dynamicSizes);
   Value zero = constantZero(rewriter, loc, elemTp);
-  rewriter.create<linalg::FillOp>(loc, zero, mem).result();
+  rewriter.create<linalg::FillOp>(loc, zero, mem);
   return mem;
 }
 
 /// Inserts the element returned by genGetNextCall(_, ind, elemPtr) into
 /// the tensor created by allocDenseTensor().  The `rank` is the rank
 /// of the `tensor` and the length of `ind`.
 static void insertScalarIntoDenseTensor(ConversionPatternRewriter &rewriter,
                                         Location loc, Value elemPtr,

 /// Sparse conversion rule for the release operator.
 class SparseTensorReleaseConverter : public OpConversionPattern<ReleaseOp> {
 public:
   using OpConversionPattern::OpConversionPattern;
   LogicalResult
   matchAndRewrite(ReleaseOp op, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {
     StringRef name = "delSparseTensor";
-    TypeRange none;
-    auto fn = getFunc(op, name, none, adaptor.getOperands());
-    rewriter.create<CallOp>(op.getLoc(), none, fn, adaptor.getOperands());
+    TypeRange noTp;
+    createFuncCall(rewriter, op, name, noTp, adaptor.getOperands());
     rewriter.eraseOp(op);
     return success();
   }
 };
 
 /// Sparse conversion rule for pointer accesses.
 class SparseTensorToPointersConverter
     : public OpConversionPattern<ToPointersOp> {
matchAndRewrite(ToPointersOp op, OpAdaptor adaptor,
     else if (eltType.isInteger(32))
       name = "sparsePointers32";
     else if (eltType.isInteger(16))
       name = "sparsePointers16";
     else if (eltType.isInteger(8))
       name = "sparsePointers8";
     else
       return failure();
-    auto fn = getFunc(op, name, resType, adaptor.getOperands(),
+    replaceOpWithFuncCall(rewriter, op, name, resType, adaptor.getOperands(),
                           /*emitCInterface=*/true);
-    rewriter.replaceOpWithNewOp<CallOp>(op, resType, fn, adaptor.getOperands());
     return success();
   }
 };
 
 /// Sparse conversion rule for index accesses.
 class SparseTensorToIndicesConverter : public OpConversionPattern<ToIndicesOp> {
 public:
   using OpConversionPattern::OpConversionPattern;
matchAndRewrite(ToIndicesOp op, OpAdaptor adaptor,
     else if (eltType.isInteger(32))
       name = "sparseIndices32";
     else if (eltType.isInteger(16))
       name = "sparseIndices16";
     else if (eltType.isInteger(8))
       name = "sparseIndices8";
     else
       return failure();
-    auto fn = getFunc(op, name, resType, adaptor.getOperands(),
+    replaceOpWithFuncCall(rewriter, op, name, resType, adaptor.getOperands(),
                           /*emitCInterface=*/true);
-    rewriter.replaceOpWithNewOp<CallOp>(op, resType, fn, adaptor.getOperands());
     return success();
   }
 };
 
 /// Sparse conversion rule for value accesses.
 class SparseTensorToValuesConverter : public OpConversionPattern<ToValuesOp> {
 public:
   using OpConversionPattern::OpConversionPattern;
matchAndRewrite(ToValuesOp op, OpAdaptor adaptor,
     else if (eltType.isInteger(32))
       name = "sparseValuesI32";
     else if (eltType.isInteger(16))
       name = "sparseValuesI16";
     else if (eltType.isInteger(8))
       name = "sparseValuesI8";
     else
       return failure();
-    auto fn = getFunc(op, name, resType, adaptor.getOperands(),
+    replaceOpWithFuncCall(rewriter, op, name, resType, adaptor.getOperands(),
                           /*emitCInterface=*/true);
-    rewriter.replaceOpWithNewOp<CallOp>(op, resType, fn, adaptor.getOperands());
     return success();
   }
 };
 
 /// Sparse conversion rule for tensor rematerialization.
 class SparseTensorLoadConverter : public OpConversionPattern<LoadOp> {
 public:
   using OpConversionPattern::OpConversionPattern;
   LogicalResult
   matchAndRewrite(LoadOp op, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {
     if (op.hasInserts()) {
       // Finalize any pending insertions.
       StringRef name = "endInsert";
       TypeRange noTp;
-      auto fn = getFunc(op, name, noTp, adaptor.getOperands());
-      rewriter.create<CallOp>(op.getLoc(), noTp, fn, adaptor.getOperands());
+      createFuncCall(rewriter, op, name, noTp, adaptor.getOperands());
     }
     rewriter.replaceOp(op, adaptor.getOperands());
     return success();
   }
 };
 
 /// Sparse conversion rule for inserting in lexicographic index order.
 class SparseTensorLexInsertConverter : public OpConversionPattern<LexInsertOp> {
else if (eltType.isInteger(32))
       name = "lexInsertI32";
     else if (eltType.isInteger(16))
       name = "lexInsertI16";
     else if (eltType.isInteger(8))
       name = "lexInsertI8";
     else
       llvm_unreachable("Unknown element type");
     TypeRange noTp;
-    auto fn =
-        getFunc(op, name, noTp, adaptor.getOperands(), /*emitCInterface=*/true);
-    rewriter.replaceOpWithNewOp<CallOp>(op, noTp, fn, adaptor.getOperands());
+    replaceOpWithFuncCall(rewriter, op, name, noTp, adaptor.getOperands(),
+                          /*emitCInterface=*/true);
     return success();
   }
 };
 
 } // namespace
 
 //===----------------------------------------------------------------------===//
 // Public method for populating conversion rules.