[mlir][sparse] Factoring out helper functions for generating constants

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

getDimLevelTypeEncoding(SparseTensorEncodingAttr::DimLevelType dlt) {
   case SparseTensorEncodingAttr::DimLevelType::Compressed:
     return 1;
   case SparseTensorEncodingAttr::DimLevelType::Singleton:
     return 2;
   }
   llvm_unreachable("Unknown SparseTensorEncodingAttr::DimLevelType");
 }
 
+/// Generates a constant zero of the given type.
+inline static Value constantZero(ConversionPatternRewriter &rewriter,
+                                 Location loc, Type t) {
+  return rewriter.create<arith::ConstantOp>(loc, t, rewriter.getZeroAttr(t));
+}
+
+/// Generates a constant of `index` type.
+inline static Value constantIndex(ConversionPatternRewriter &rewriter,
+                                  Location loc, unsigned i) {
+  return rewriter.create<arith::ConstantIndexOp>(loc, i);
+}
+
+/// Generates a constant of `i64` type.
+inline static Value constantI64(ConversionPatternRewriter &rewriter,
+                                Location loc, int64_t i) {
+  return rewriter.create<arith::ConstantIntOp>(loc, i, 64);
+}
+
+/// Generates a constant of `i32` type.
+inline static Value constantI32(ConversionPatternRewriter &rewriter,
+                                Location loc, int32_t i) {
+  return rewriter.create<arith::ConstantIntOp>(loc, i, 32);
+}
+
 /// Returns integers of given width and values as a constant tensor.
 /// We cast the static shape into a dynamic shape to ensure that the
 /// method signature remains uniform across different tensor dimensions.
 static Value getTensor(ConversionPatternRewriter &rewriter, unsigned width,
                        Location loc, ArrayRef<APInt> values) {
   Type etp = rewriter.getIntegerType(width);
   unsigned sz = values.size();
   RankedTensorType tt1 = RankedTensorType::get({sz}, etp);
static Value genNewCall(ConversionPatternRewriter &rewriter, Operation *op,
   }
   perm = getTensor(rewriter, 64, loc, rev);
   params.push_back(perm);
   // Secondary and primary types encoding.
   unsigned secPtr = getOverheadTypeEncoding(enc.getPointerBitWidth());
   unsigned secInd = getOverheadTypeEncoding(enc.getIndexBitWidth());
   unsigned primary = getPrimaryTypeEncoding(resType.getElementType());
   assert(primary);
-  params.push_back(rewriter.create<arith::ConstantOp>(
-      loc, rewriter.getI64IntegerAttr(secPtr)));
-  params.push_back(rewriter.create<arith::ConstantOp>(
-      loc, rewriter.getI64IntegerAttr(secInd)));
-  params.push_back(rewriter.create<arith::ConstantOp>(
-      loc, rewriter.getI64IntegerAttr(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));
   if (!ptr)
     ptr = rewriter.create<LLVM::NullOp>(loc, pTp);
-  params.push_back(rewriter.create<arith::ConstantOp>(
-      loc, rewriter.getI32IntegerAttr(action)));
+  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),
       params);
   return call.getResult(0);
 }
 
-/// Generates a constant zero of the given type.
-static Value getZero(ConversionPatternRewriter &rewriter, Location loc,
-                     Type t) {
-  return rewriter.create<arith::ConstantOp>(loc, rewriter.getZeroAttr(t));
-}
-
 /// 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) {
   Type t = v.getType();
-  Value zero = getZero(rewriter, loc, t);
+  Value zero = constantZero(rewriter, loc, t);
   if (t.isa<FloatType>())
     return rewriter.create<arith::CmpFOp>(loc, arith::CmpFPredicate::UNE, v,
                                           zero);
   if (t.isIntOrIndex())
     return rewriter.create<arith::CmpIOp>(loc, arith::CmpIPredicate::ne, v,
                                           zero);
   llvm_unreachable("Unknown element type");
 }
static Value genIndexAndValueForDense(ConversionPatternRewriter &rewriter,
                                       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 =
-        rewriter.create<arith::ConstantOp>(loc, rewriter.getIndexAttr(i++));
+    Value idx = constantIndex(rewriter, loc, i++);
     rewriter.create<memref::StoreOp>(loc, iv, ind, idx);
   }
   return val;
 }
 
 /// Generates a call that adds one element to a coordinate scheme.
 /// In particular, this generates code like the following:
 ///   val = a[i1,..,ik];

 /// 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,
                                        Value values, Value ind, ValueRange ivs,
                                        unsigned rank) {
   Location loc = op->getLoc();
   for (unsigned i = 0; i < rank; i++) {
-    Value idx =
-        rewriter.create<arith::ConstantOp>(loc, rewriter.getIndexAttr(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);
   }
   return rewriter.create<tensor::ExtractOp>(loc, values, ivs[0]);
 }
 
 /// Generates code to stack-allocate a `memref<?xindex>` where the `?`
 /// is the given `rank`.  This array is intended to serve as a reusable
 /// buffer for storing the indices of a single tensor element, to avoid
 /// allocation in the body of loops.
 static Value allocaIndices(ConversionPatternRewriter &rewriter, Location loc,
                            int64_t rank) {
   auto indexTp = rewriter.getIndexType();
   auto memTp = MemRefType::get({ShapedType::kDynamicSize}, indexTp);
-  Value arg =
-      rewriter.create<arith::ConstantOp>(loc, rewriter.getIndexAttr(rank));
+  Value arg = constantIndex(rewriter, loc, rank);
   return rewriter.create<memref::AllocaOp>(loc, memTp, ValueRange{arg});
 }
 
 //===----------------------------------------------------------------------===//
 // Conversion rules.
 //===----------------------------------------------------------------------===//
 
 /// Sparse conversion rule for returns.
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<Value, 2> params;
     params.push_back(adaptor.getOperands()[0]);
-    params.push_back(rewriter.create<arith::ConstantOp>(
-        op.getLoc(), rewriter.getIndexAttr(idx)));
+    params.push_back(constantIndex(rewriter, op.getLoc(), idx));
     rewriter.replaceOpWithNewOp<CallOp>(
         op, resType, getFunc(op, name, resType, params), params);
     return success();
   }
 };
 
 /// Sparse conversion rule for the new operator.
 class SparseTensorNewConverter : public OpConversionPattern<NewOp> {
matchAndRewrite(ConvertOp op, OpAdaptor adaptor,
     Location loc = op->getLoc();
     ShapedType shape = resType.cast<ShapedType>();
     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 =
-        rewriter.create<arith::ConstantOp>(loc, rewriter.getIndexAttr(0));
-    Value one =
-        rewriter.create<arith::ConstantOp>(loc, rewriter.getIndexAttr(1));
+    Value zero = constantIndex(rewriter, loc, 0);
+    Value one = constantIndex(rewriter, loc, 1);
     auto indicesValues = genSplitSparseConstant(rewriter, op, src);
     bool isCOOConstant = indicesValues.hasValue();
     Value indices;
     Value values;
     if (isCOOConstant) {
       indices = indicesValues->first;
       values = indicesValues->second;
       lo.push_back(zero);