[mlir][linalg] Cleanup LinalgOp usage in sparse compiler (NFC).

[gysit]

Summary

Replace the uses of deprecated Structured Op Interface methods in Sparsification.cpp. This patch is based on https://reviews.llvm.org/D103394.

[aartbik]

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

567

why is this hoisted so far before its use?

[gysit]

Address comment.

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

static Dim toDim(SparseTensorEncodingAttr &enc, unsigned d) {
   }
   return Dim::kDense;
 }
 
 /// Helper method to inspect sparse encodings in the tensor types.
 /// Fills the per-dimension sparsity information for all tensors.
 static bool findSparseAnnotations(Merger &merger, linalg::GenericOp op) {
   bool annotated = false;
-  unsigned numTensors = op.getNumShapedOperands();
-  unsigned lhs = numTensors - 1;
-  for (unsigned t = 0; t < numTensors; t++) {
-    auto map = op.getIndexingMap(t);
+  OpOperand *lhs = op.getOutputOperand(0);
+  for (OpOperand *t : op.getInputAndOutputOperands()) {
+    auto map = op.getTiedIndexingMap(t);
     if (!map.isProjectedPermutation())
       return false;
-    auto enc = getSparseTensorEncoding(op.getShapedType(t));
+    auto enc = getSparseTensorEncoding(t->get().getType());
     if (enc) {
       annotated = true;
       if (t == lhs)
         return false; // TODO: handle sparse outputs
     }
-    assert(map.getNumResults() == op.getShapedType(t).getRank());
+    assert(map.getNumResults() == op.getRank(t));
     for (unsigned d = 0, rank = map.getNumResults(); d < rank; d++) {
       unsigned idx = map.getDimPosition(perm(enc, d));
-      merger.setDim(t, idx, toDim(enc, d));
+      merger.setDim(t->getOperandNumber(), idx, toDim(enc, d));
     }
   }
   return annotated;
 }
 
 /// A DFS helper to compute a topological sort. Note that recursion is
 /// bounded by the number of implicit loops, which is always small.
 /// Returns false when a cycle is detected.
static bool computeIterationGraph(Merger &merger, linalg::GenericOp op,
                                   std::vector<unsigned> &topSort,
                                   bool sparseOnly) {
   // Set up an n x n from/to adjacency matrix of the iteration graph
   // for the implicit loop indices i_0 .. i_n-1.
   unsigned n = op.getNumLoops();
   std::vector<std::vector<bool>> adjM(n, std::vector<bool>(n, false));
 
   // Iterate over the indexing maps of every tensor in the tensor expression.
-  unsigned numTensors = op.getNumShapedOperands();
-  for (unsigned t = 0; t < numTensors; t++) {
-    auto map = op.getIndexingMap(t);
-    auto enc = getSparseTensorEncoding(op.getShapedType(t));
+  for (OpOperand *t : op.getInputAndOutputOperands()) {
+    auto map = op.getTiedIndexingMap(t);
+    auto enc = getSparseTensorEncoding(t->get().getType());
     assert(map.getNumDims() == n);
     // Skip dense tensor constraints when sparse only is requested.
     if (sparseOnly && !enc)
       continue;
     // Each tensor expression and optional dimension ordering (row-major
     // by default) puts an ordering constraint on the loop indices. For
     // example, the tensor expresion A_ijk forces the ordering i < j < k
     // on the loop indices if no explicit dimension ordering is given.

 static unsigned buildLattices(Merger &merger, linalg::GenericOp op,
                               unsigned exp, unsigned idx) {
   Kind kind = merger.exp(exp).kind;
   if (kind == Kind::kTensor || kind == Kind::kInvariant) {
     // Either the index is really used in the tensor expression, or it is
     // set to the undefined index in that dimension. An invariant expression
     // is set to a synthetic tensor with undefined indices only.
     unsigned s = merger.addSet();
-    unsigned t =
-        kind == Kind::kTensor ? merger.exp(exp).e0 : op.getNumShapedOperands();
+    unsigned t = kind == Kind::kTensor ? merger.exp(exp).e0
+                                       : op.getNumInputsAndOutputs();
     merger.set(s).push_back(merger.addLat(t, idx, exp));
     return s;
   }
   unsigned s0 = buildLattices(merger, op, merger.exp(exp).e0, idx);
   unsigned s1 = buildLattices(merger, op, merger.exp(exp).e1, idx);
   switch (kind) {
   case Kind::kTensor:
   case Kind::kInvariant:
static bool getInPlace(Value val) {
   return false;
 }
 
 /// Generates buffer for the output tensor.
 static Value genOutputBuffer(CodeGen &codegen, PatternRewriter &rewriter,
                              linalg::GenericOp op, MemRefType denseTp,
                              ArrayRef<Value> args) {
   Location loc = op.getLoc();
-  Value tensor = op.getOutput(0);
+  Value tensor = op.getOutputOperand(0)->get();
   // The output tensor simply could materialize from the buffer that will
   // be generated for the tensor present in the outs() clause. This has
   // the major advantage that the sparse kernel only updates the nonzero
   // positions for the output tensor.
   if (getInPlace(tensor))
     return rewriter.create<memref::BufferCastOp>(loc, denseTp, tensor);
   // By default, a new buffer is allocated which is initialized to the
   // tensor defined in the outs() clause. This is always correct but
   // introduces a dense initialization component that may negatively
   // impact the running complexity of the sparse kernel.
   Value init = rewriter.create<memref::BufferCastOp>(loc, denseTp, tensor);
   Value alloc = rewriter.create<memref::AllocOp>(loc, denseTp, args);
   rewriter.create<linalg::CopyOp>(loc, init, alloc);
   return alloc;
 }
 
 /// Local bufferization of all dense and sparse data structures.
 /// This code enables testing the first prototype sparse compiler.
 // TODO: replace this with a proliferated bufferization strategy
 static void genBuffers(Merger &merger, CodeGen &codegen,
                        PatternRewriter &rewriter, linalg::GenericOp op) {
   Location loc = op.getLoc();
-  unsigned numTensors = op.getNumShapedOperands();
-  unsigned numInputs = op.getNumInputs();
-  assert(numTensors == numInputs + 1);
+  assert(op.getNumInputsAndOutputs() == op.getNumInputs() + 1);
   // For every tensor, find lower and upper bound on dimensions, set the
   // same bounds on loop indices, and obtain dense or sparse buffer(s).
   SmallVector<Value, 4> args;
-  for (unsigned t = 0; t < numTensors; t++) {
-    Value tensor = t < numInputs ? op.getInput(t) : op.getOutput(0);
-    auto tensorType = op.getShapedType(t);
-    auto shape = tensorType.getShape();
-    auto map = op.getIndexingMap(t);
-    auto enc = getSparseTensorEncoding(tensorType);
+  for (OpOperand *t : op.getInputAndOutputOperands()) {
+    auto shape = op.getShape(t);
+    auto map = op.getTiedIndexingMap(t);
+    auto enc = getSparseTensorEncoding(t->get().getType());
     // Scan all dimensions of current tensor.
     args.clear();
     for (unsigned d = 0, rank = map.getNumResults(); d < rank; d++) {
       unsigned idx = map.getDimPosition(perm(enc, d));
       // Handle sparse storage schemes.
-      if (merger.isDim(t, idx, Dim::kSparse)) {
+      if (merger.isDim(t->getOperandNumber(), idx, Dim::kSparse)) {
         auto dynShape = {ShapedType::kDynamicSize};
         auto ptrTp = MemRefType::get(
             dynShape, genIntType(rewriter, enc.getPointerBitWidth()));
         auto indTp = MemRefType::get(
             dynShape, genIntType(rewriter, enc.getIndexBitWidth()));
         Value dim = rewriter.create<ConstantIndexOp>(loc, d);
         // Generate sparse primitives to obtains pointer and indices.
-        codegen.pointers[t][idx] =
-            rewriter.create<ToPointersOp>(loc, ptrTp, tensor, dim);
-        codegen.indices[t][idx] =
-            rewriter.create<ToIndicesOp>(loc, indTp, tensor, dim);
+        codegen.pointers[t->getOperandNumber()][idx] =
+            rewriter.create<ToPointersOp>(loc, ptrTp, t->get(), dim);
+        codegen.indices[t->getOperandNumber()][idx] =
+            rewriter.create<ToIndicesOp>(loc, indTp, t->get(), dim);
       }
       // Find lower and upper bound in current dimension.
       Value up;
       if (shape[d] == MemRefType::kDynamicSize) {
-        up = rewriter.create<memref::DimOp>(loc, tensor, d);
+        up = rewriter.create<memref::DimOp>(loc, t->get(), d);
         args.push_back(up);
       } else {
         up = rewriter.create<ConstantIndexOp>(loc, shape[d]);
       }
-      codegen.sizes[idx] = codegen.highs[t][idx] = up;
+      codegen.sizes[idx] = codegen.highs[t->getOperandNumber()][idx] = up;
     }
     // Perform the required bufferization. All dense inputs materialize
     // from the input tensor. The dense output tensor needs special
     // handling. Sparse inputs use a sparse primitive to obtain the values.
+    Type elementType = getElementTypeOrSelf(t->get().getType());
     if (!enc) {
-      auto denseTp = MemRefType::get(shape, tensorType.getElementType());
-      if (t < numInputs)
-        codegen.buffers[t] =
-            rewriter.create<memref::BufferCastOp>(loc, denseTp, tensor);
+      auto denseTp = MemRefType::get(shape, elementType);
+      if (t->getOperandNumber() < op.getNumInputs())
+        codegen.buffers[t->getOperandNumber()] =
+            rewriter.create<memref::BufferCastOp>(loc, denseTp, t->get());
       else
-        codegen.buffers[t] =
+        codegen.buffers[t->getOperandNumber()] =
             genOutputBuffer(codegen, rewriter, op, denseTp, args);
     } else {
       auto dynShape = {ShapedType::kDynamicSize};
-      auto sparseTp = MemRefType::get(dynShape, tensorType.getElementType());
-      codegen.buffers[t] = rewriter.create<ToValuesOp>(loc, sparseTp, tensor);
+      auto sparseTp = MemRefType::get(dynShape, elementType);
+      codegen.buffers[t->getOperandNumber()] =
+          rewriter.create<ToValuesOp>(loc, sparseTp, t->get());
     }
   }
 }
 
 /// Constructs vector type.
 static VectorType vectorType(CodeGen &codegen, Type etp) {
   return VectorType::get(codegen.curVecLength, etp);
 }
static Value genTensorLoad(Merger &merger, CodeGen &codegen,
   Value val = merger.exp(exp).val;
   if (val) {
     if (codegen.curVecLength > 1 && !val.getType().isa<VectorType>())
       return genVectorInvariantValue(codegen, rewriter, val);
     return val;
   }
   // Actual load.
   SmallVector<Value, 4> args;
-  unsigned tensor = merger.exp(exp).e0;
-  auto map = op.getIndexingMap(tensor);
-  auto enc = getSparseTensorEncoding(op.getShapedType(tensor));
+  OpOperand *tensor = merger.exp(exp).e0 < op.getNumInputs()
+                          ? op.getInputOperand(merger.exp(exp).e0)
+                          : op.getOutputOperand(0);
+  auto map = op.getTiedIndexingMap(tensor);
+  auto enc = getSparseTensorEncoding(tensor->get().getType());
   for (unsigned d = 0, rank = map.getNumResults(); d < rank; d++) {
     unsigned idx = map.getDimPosition(perm(enc, d));
     args.push_back(codegen.loops[idx]); // universal dense index
     if (enc) {
       args.clear();
-      args.push_back(codegen.pidxs[tensor][idx]); // position index
+      args.push_back(
+          codegen.pidxs[tensor->getOperandNumber()][idx]); // position index
     }
   }
   Location loc = op.getLoc();
-  Value ptr = codegen.buffers[tensor];
+  Value ptr = codegen.buffers[tensor->getOperandNumber()];
   if (codegen.curVecLength > 1)
     return genVectorLoad(codegen, rewriter, ptr, args);
   return rewriter.create<memref::LoadOp>(loc, ptr, args);
 }
 
 /// Generates a store on a dense tensor.
 static void genTensorStore(Merger &merger, CodeGen &codegen,
                            PatternRewriter &rewriter, linalg::GenericOp op,
-                           unsigned tensor, Value rhs) {
+                           OpOperand *tensor, Value rhs) {
   Location loc = op.getLoc();
   // Test if this is a scalarized reduction.
-  unsigned lhs = op.getNumShapedOperands() - 1;
+  OpOperand *lhs = op.getOutputOperand(0);
   if (lhs == tensor && codegen.redVal) {
     if (codegen.curVecLength > 1)
       rhs = rewriter.create<SelectOp>(loc, codegen.curVecMask, rhs,
                                       codegen.redVal);
     codegen.redVal = rhs;
     return;
   }
   // Actual store.
   SmallVector<Value, 4> args;
-  auto map = op.getIndexingMap(tensor);
-  assert(!getSparseTensorEncoding(op.getShapedType(tensor)));
+  auto map = op.getTiedIndexingMap(tensor);
+  assert(!getSparseTensorEncoding(tensor->get().getType()));
   for (unsigned d = 0, rank = map.getNumResults(); d < rank; d++) {
     unsigned idx = map.getDimPosition(d);
     args.push_back(codegen.loops[idx]); // universal dense index
   }
-  Value ptr = codegen.buffers[tensor];
+  Value ptr = codegen.buffers[tensor->getOperandNumber()];
   if (codegen.curVecLength > 1)
     genVectorStore(codegen, rewriter, rhs, ptr, args);
   else
     rewriter.create<memref::StoreOp>(loc, rhs, ptr, args);
 }
 
 /// Generates a pointer/index load from the sparse storage scheme. Narrower
 /// data types need to be zero extended before casting the value into the

 /// Generates end of a reduction.
 static void genReductionEnd(Merger &merger, CodeGen &codegen,
                             PatternRewriter &rewriter, linalg::GenericOp op) {
   Value red = codegen.redVal;
   if (!red)
     return;
   assert(codegen.curVecLength == 1);
   codegen.redVal = merger.exp(codegen.redExp).val = Value(); // end chain
-  unsigned lhs = op.getNumShapedOperands() - 1;
+  OpOperand *lhs = op.getOutputOperand(0);
   if (auto vtp = red.getType().dyn_cast<VectorType>()) {
     // TODO: assumes + reductions for now
     StringAttr kind = rewriter.getStringAttr("add");
     Value ld = genTensorLoad(merger, codegen, rewriter, op, codegen.redExp);
     // Integer reductions don't accept an accumulator.
     if (vtp.getElementType().isa<IntegerType>()) {
       red = rewriter.create<vector::ReductionOp>(op.getLoc(), ld.getType(),
                                                  kind, red, ValueRange{});

 
 /// Hoists loop invariant tensor loads for which indices have been exhausted.
 static void genInvariants(Merger &merger, CodeGen &codegen,
                           PatternRewriter &rewriter, linalg::GenericOp op,
                           unsigned exp, unsigned ldx, bool hoist) {
   if (merger.exp(exp).kind == Kind::kTensor) {
     // Inspect tensor indices.
     bool atLevel = ldx == -1u;
-    unsigned tensor = merger.exp(exp).e0;
-    auto map = op.getIndexingMap(tensor);
-    auto enc = getSparseTensorEncoding(op.getShapedType(tensor));
+    OpOperand *tensor = merger.exp(exp).e0 < op.getNumInputs()
+                            ? op.getInputOperand(merger.exp(exp).e0)
+                            : op.getOutputOperand(0);
+    auto map = op.getTiedIndexingMap(tensor);
+    auto enc = getSparseTensorEncoding(tensor->get().getType());
     for (unsigned d = 0, rank = map.getNumResults(); d < rank; d++) {
       unsigned idx = map.getDimPosition(perm(enc, d));
       if (!codegen.loops[idx])
         return; // still in play
       else if (idx == ldx)
         atLevel = true;
     }
     // All exhausted at this level (atLevel denotes exactly at this level).
-    unsigned lhs = op.getNumShapedOperands() - 1;
+    OpOperand *lhs = op.getOutputOperand(0);
     if (lhs == tensor) {
       codegen.redExp = hoist ? exp : -1u;
     } else if (atLevel) {
       merger.exp(exp).val =
           hoist ? genTensorLoad(merger, codegen, rewriter, op, exp) : Value();
     }
   } else if (merger.exp(exp).kind != Kind::kInvariant) {
     // Traverse into the binary operations. Note that we only hoist

 
 /// Checks unit strides for dense tensors. The iteration graph may have ignored
 /// dense access patterns in order to avoid cycles (sparse access patterns are
 /// always placed innermost), but that means dense access has become strided.
 /// For now, we reject vectorization of such cases.
 /// TODO: implement strided load/stores on dense arrays
 static bool denseUnitStrides(Merger &merger, linalg::GenericOp op,
                              unsigned idx) {
-  unsigned numTensors = op.getNumShapedOperands();
-  for (unsigned t = 0; t < numTensors; t++) {
-    if (!getSparseTensorEncoding(op.getShapedType(t))) {
-      auto map = op.getIndexingMap(t);
+  for (OpOperand *t : op.getInputAndOutputOperands()) {
+    if (!getSparseTensorEncoding(t->get().getType())) {
+      auto map = op.getTiedIndexingMap(t);
       for (unsigned d = 0, rank = map.getNumResults(); d < rank; d++) {
         if (map.getDimPosition(d) == idx && d != rank - 1)
           return false;
       }
     }
   }
   return true;
 }

 /// Recursively generates code while computing iteration lattices in order
 /// to manage the complexity of implementing co-iteration over unions
 /// and intersections of sparse iterations spaces.
 static void genStmt(Merger &merger, CodeGen &codegen, PatternRewriter &rewriter,
                     linalg::GenericOp op, std::vector<unsigned> &topSort,
                     unsigned exp, unsigned at) {
   // At each leaf, assign remaining tensor (sub)expression to output tensor.
   if (at == topSort.size()) {
-    unsigned lhs = op.getNumShapedOperands() - 1;
+    OpOperand *lhs = op.getOutputOperand(0);
     Value rhs = genExp(merger, codegen, rewriter, op, exp);
     genTensorStore(merger, codegen, rewriter, op, lhs, rhs);
     return;
   }
   assert(codegen.curVecLength == 1);
 
   // Construct iteration lattices for current loop index, with L0 at top.
   // Then emit initialization code for the loop sequence at this level.
public:
   GenericOpSparsifier(MLIRContext *context, SparsificationOptions o)
       : OpRewritePattern<linalg::GenericOp>(context), options(o) {}
 
   LogicalResult matchAndRewrite(linalg::GenericOp op,
                                 PatternRewriter &rewriter) const override {
     // Detects sparse annotations and translate the per-dimension sparsity
     // information for all tensors to loop indices in the kernel.
     assert(op.getNumOutputs() == 1);
-    unsigned numTensors = op.getNumShapedOperands();
+    unsigned numTensors = op.getNumInputsAndOutputs();
     unsigned numLoops = op.iterator_types().getValue().size();
     Merger merger(numTensors, numLoops);
     if (!findSparseAnnotations(merger, op))
       return failure();
 
     // Computes a topologically sorted iteration graph to ensure
     // tensors are visited in natural index order. Fails on cycles.
     // This assumes that higher-level passes have already put the