[mlir][sparse] add some APIs for merger to query the tensor id for output tensor and synthetic tensor.

[aartbik]

mlir/include/mlir/Dialect/SparseTensor/Utils/Merger.h

233

I think the /// comment technically needs to be "per" method, not for a group, so

/// Gets the id of the output tensor.
unsigned getOutTensorID() const { return outTensor; }

/// Gets the id of the synthetic tensor (used for all invariant tensor expressions).
unsigned getSynTensorID() const { return syntheticTensor; }

[Peiming]

address comments.

mlir/include/mlir/Dialect/SparseTensor/Utils/Merger.h

public:
   unsigned tensor(unsigned b) const { return b % numTensors; }
   unsigned index(unsigned b) const { return b / numTensors; }
 
   /// Returns true if bit corresponds to index of output tensor.
   bool isOutTensor(unsigned b, unsigned i) const {
     return tensor(b) == outTensor && index(b) == i;
   }
 
+  /// Gets tensor ID for special tensors
+  unsigned getOutTensorID() const { return outTensor; }
+  unsigned getSynTensorID() const { return syntheticTensor; }
+
   /// Returns true if given tensor iterates *only* in the given tensor
   /// expression. For the output tensor, this defines a "simply dynamic"
   /// operation [Bik96]. For instance: a(i) *= 2.0 or a(i) += a(i) for
   /// sparse vector a.
   bool isSingleCondition(unsigned t, unsigned e) const;
 
   /// Returns true if any set bit corresponds to sparse dimension level type.
   bool hasAnySparse(const BitVector &bits) const;

mlir/unittests/Dialect/SparseTensor/MergerTest.cpp

 protected:
   // Our three tensors (two inputs, one output).
   const unsigned t0 = 0, t1 = 1, t2 = 2;
 
   // Our single loop.
   const unsigned l0 = 0;
 
   MergerTest3T1L() : MergerTestBase(3, 1) {
+    EXPECT_TRUE(merger.getOutTensorID() == t2);
+
     // Tensor 0: sparse input vector.
     merger.addExp(Kind::kTensor, t0, -1u);
     merger.setDimLevelType(t0, l0, DimLevelType::Compressed);
 
     // Tensor 1: sparse input vector.
     merger.addExp(Kind::kTensor, t1, -1u);
     merger.setDimLevelType(t1, l0, DimLevelType::Compressed);
 
     // Tensor 2: dense output vector.
     merger.addExp(Kind::kTensor, t2, -1u);
     merger.setDimLevelType(t2, l0, DimLevelType::Dense);
   }
 };
 
 class MergerTest4T1L : public MergerTestBase {
 protected:
   // Our four tensors (three inputs, one output).
   const unsigned t0 = 0, t1 = 1, t2 = 2, t3 = 3;
 
   // Our single loop.
   const unsigned l0 = 0;
 
   MergerTest4T1L() : MergerTestBase(4, 1) {
+    EXPECT_TRUE(merger.getOutTensorID() == t3);
+
     // Tensor 0: sparse input vector.
     merger.addExp(Kind::kTensor, t0, -1u);
     merger.setDimLevelType(t0, l0, DimLevelType::Compressed);
 
     // Tensor 1: sparse input vector.
     merger.addExp(Kind::kTensor, t1, -1u);
     merger.setDimLevelType(t1, l0, DimLevelType::Compressed);
 

 protected:
   // Our three tensors (two inputs, one output).
   const unsigned t0 = 0, t1 = 1, t2 = 2;
 
   // Our single loop.
   const unsigned l0 = 0;
 
   MergerTest3T1LD() : MergerTestBase(3, 1) {
+    EXPECT_TRUE(merger.getOutTensorID() == t2);
+
     // Tensor 0: sparse input vector.
     merger.addExp(Kind::kTensor, t0, -1u);
     merger.setDimLevelType(t0, l0, DimLevelType::Compressed);
 
     // Tensor 1: dense input vector.
     merger.addExp(Kind::kTensor, t1, -1u);
     merger.setDimLevelType(t1, l0, DimLevelType::Dense);
 
     // Tensor 2: dense output vector.
     merger.addExp(Kind::kTensor, t2, -1u);
     merger.setDimLevelType(t2, l0, DimLevelType::Dense);
   }
 };
 
 ///
 /// Tests with both undef and dense input.
 ///
 
 class MergerTest4T1LU : public MergerTestBase {
 protected:
-  // Our three tensors (two inputs, one output).
+  // Our three tensors (three inputs, one output).
   const unsigned t0 = 0, t1 = 1, t2 = 2, t3 = 3;
 
   // Our single loop.
   const unsigned l0 = 0;
 
   MergerTest4T1LU() : MergerTestBase(4, 1) {
+    EXPECT_TRUE(merger.getOutTensorID() == t3);
+
     // Tensor 0: undef input vector.
     merger.addExp(Kind::kTensor, t0, -1u);
     merger.setDimLevelType(t0, l0, DimLevelType::Undef);
 
     // Tensor 1: dense input vector.
     merger.addExp(Kind::kTensor, t1, -1u);
     merger.setDimLevelType(t1, l0, DimLevelType::Dense);
 

 protected:
   // Our three tensors (two inputs, one output, one synthetic).
   const unsigned t0 = 0, t1 = 1, t2 = 2, t3 = 3;
 
   // Our single loop.
   const unsigned l0 = 0;
 
   MergerTest3T1L_SO() : MergerTestBase(3, 1) {
+    EXPECT_TRUE(merger.getOutTensorID() == t2);
+    EXPECT_TRUE(merger.getSynTensorID() == t3);
+
     merger.setHasSparseOut(true);
 
     // Tensor 0: undef input vector.
     merger.addExp(Kind::kTensor, t0, -1u);
     merger.setDimLevelType(t0, l0, DimLevelType::Undef);
 
     // Tensor 1: undef input vector.
     merger.addExp(Kind::kTensor, t1, -1u);