[mlir][sparse] Misc code cleanup

[wrengr]

Summary

Depends On D111763

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

static Value genNewCall(ConversionPatternRewriter &rewriter, Operation *op,
   unsigned secPtr = getOverheadTypeEncoding(enc.getPointerBitWidth());
   unsigned secInd = getOverheadTypeEncoding(enc.getIndexBitWidth());
   unsigned primary = getPrimaryTypeEncoding(resType.getElementType());
   assert(primary);
   params.push_back(constantI64(rewriter, loc, secPtr));
   params.push_back(constantI64(rewriter, loc, secInd));
   params.push_back(constantI64(rewriter, loc, primary));
   // User action and pointer.
-  Type pTp = LLVM::LLVMPointerType::get(IntegerType::get(op->getContext(), 8));
+  Type pTp = LLVM::LLVMPointerType::get(rewriter.getI8Type());
   if (!ptr)
     ptr = rewriter.create<LLVM::NullOp>(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<CallOp>(
       loc, pTp, getFunc(op, name, pTp, params, /*emitCInterface=*/true),

 /// 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(ConversionPatternRewriter &rewriter,
-                                      Operation *op, Value tensor, Value ind,
+                                      Location loc, Value tensor, Value ind,
                                       ValueRange ivs) {
-  Location loc = op->getLoc();
   Value val = rewriter.create<tensor::ExtractOp>(loc, tensor, ivs);
   Value cond = genIsNonzero(rewriter, loc, val);
   scf::IfOp ifOp = rewriter.create<scf::IfOp>(loc, cond, /*else*/ false);
   rewriter.setInsertionPointToStart(&ifOp.thenRegion().front());
   unsigned i = 0;
   for (auto iv : ivs) {
     Value idx = constantIndex(rewriter, loc, i++);
     rewriter.create<memref::StoreOp>(loc, iv, ind, idx);
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(IntegerType::get(op->getContext(), 8));
+  Type pTp = LLVM::LLVMPointerType::get(rewriter.getI8Type());
   rewriter.create<CallOp>(
       loc, pTp, getFunc(op, name, pTp, params, /*emitCInterface=*/true),
       params);
 }
 
 /// 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, ConvertOp op,
+genSplitSparseConstant(ConversionPatternRewriter &rewriter, Location loc,
                        Value tensor) {
   if (auto constOp = tensor.getDefiningOp<arith::ConstantOp>()) {
     if (auto attr = constOp.value().dyn_cast<SparseElementsAttr>()) {
-      Location loc = op->getLoc();
       DenseElementsAttr indicesAttr = attr.getIndices();
       Value indices = rewriter.create<arith::ConstantOp>(loc, indicesAttr);
       DenseElementsAttr valuesAttr = attr.getValues();
       Value values = rewriter.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(ConversionPatternRewriter &rewriter,
-                                       Operation *op, Value indices,
+                                       Location loc, Value indices,
                                        Value values, Value ind, ValueRange ivs,
                                        unsigned rank) {
-  Location loc = op->getLoc();
   for (unsigned i = 0; i < rank; i++) {
     Value idx = constantIndex(rewriter, loc, i);
     Value val = rewriter.create<tensor::ExtractOp>(loc, indices,
                                                    ValueRange{ivs[0], idx});
     val =
         rewriter.create<arith::IndexCastOp>(loc, val, rewriter.getIndexType());
     rewriter.create<memref::StoreOp>(loc, val, ind, idx);
   }
matchAndRewrite(ConvertOp op, OpAdaptor adaptor,
     Value perm;
     Value ptr = genNewCall(rewriter, op, encDst, 2, perm);
     Value ind = allocaIndices(rewriter, loc, shape.getRank());
     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, op, src);
+    auto indicesValues = genSplitSparseConstant(rewriter, loc, src);
     bool isCOOConstant = indicesValues.hasValue();
     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, rank = shape.getRank(); 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();
     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, op, indices, values, ind,
+            val = genIndexAndValueForSparse(rewriter, loc, indices, values, ind,
                                             ivs, rank);
           else
-            val = genIndexAndValueForDense(rewriter, op, src, ind, ivs);
+            val = genIndexAndValueForDense(rewriter, loc, src, ind, ivs);
           genAddEltCall(rewriter, op, eltType, ptr, val, ind, perm);
           return {};
         });
     rewriter.replaceOp(op, genNewCall(rewriter, op, encDst, 1, perm, ptr));
     return success();
   }
 };