[mlir][sparse] refactoring: split translateIndices.

[Peiming]

Summary

TranslateIndicesArray take an array of SSA value and convert them into another array of SSA values based on reassociation. Which makes it easier to be reused by foreach operator (as the indices array are given as an array of SSA values).

[Peiming]

fix name

[aartbik]

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

242

collapsion -> collapse

[Peiming]

address comments.

[bixia]

mlir/lib/Dialect/SparseTensor/Transforms/CodegenUtils.h

198

We usually put user facing doc for the routine on the header file not the cpp implementation.

[Peiming]

address comments

mlir/lib/Dialect/SparseTensor/Transforms/CodegenUtils.h

 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_DIALECT_SPARSETENSOR_TRANSFORMS_CODEGENUTILS_H_
 #define MLIR_DIALECT_SPARSETENSOR_TRANSFORMS_CODEGENUTILS_H_
 
 #include "mlir/Dialect/Arith/IR/Arith.h"
 #include "mlir/Dialect/Complex/IR/Complex.h"
 #include "mlir/Dialect/SparseTensor/IR/SparseTensor.h"
+#include "mlir/Dialect/Utils/ReshapeOpsUtils.h"
 #include "mlir/ExecutionEngine/SparseTensor/Enums.h"
+#include "mlir/ExecutionEngine/SparseTensorUtils.h"
 #include "mlir/IR/Builders.h"
 
 namespace mlir {
 class Location;
 class Type;
 class Value;
 
 namespace sparse_tensor {

 /// Generates a constant of the internal dimension level type encoding.
 inline Value
 constantDimLevelTypeEncoding(OpBuilder &builder, Location loc,
                              SparseTensorEncodingAttr::DimLevelType dlt) {
   return constantI8(builder, loc,
                     static_cast<uint8_t>(dimLevelTypeEncoding(dlt)));
 }
 
+void translateIndicesArray(OpBuilder &builder, Location loc,
+                           ArrayRef<ReassociationIndices> reassociation,
+                           ValueRange srcIndices, ArrayRef<Value> srcShape,
+                           ArrayRef<Value> dstShape,
+                           SmallVectorImpl<Value> &dstIndices);
 } // namespace sparse_tensor
 } // namespace mlir
 
 #endif // MLIR_DIALECT_SPARSETENSOR_TRANSFORMS_CODEGENUTILS_H_

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

return builder.create<arith::CmpFOp>(loc, arith::CmpFPredicate::UNE, v,
                                          zero);
   if (tp.isIntOrIndex())
     return builder.create<arith::CmpIOp>(loc, arith::CmpIPredicate::ne, v,
                                          zero);
   if (tp.dyn_cast<ComplexType>())
     return builder.create<complex::NotEqualOp>(loc, v, zero);
   llvm_unreachable("Non-numeric type");
 }
+
+/// Helper method to translate indices during a reshaping operation.
+void mlir::sparse_tensor::translateIndicesArray(
+    OpBuilder &builder, Location loc,
+    ArrayRef<ReassociationIndices> reassociation, ValueRange srcIndices,
+    ArrayRef<Value> srcShape, ArrayRef<Value> dstShape,
+    SmallVectorImpl<Value> &dstIndices) {
+  unsigned i = 0;
+  unsigned start = 0;
+  unsigned dstRank = dstShape.size();
+  unsigned srcRank = srcShape.size();
+  assert(srcRank == srcIndices.size());
+  bool isExpand = srcRank > dstRank;
+  ArrayRef<Value> shape = isExpand ? srcShape : dstShape;
+  // Iterate over reassociation map.
+  for (const auto &map : llvm::enumerate(reassociation)) {
+    // Prepare strides information in dimension slice.
+    Value linear = constantIndex(builder, loc, 1);
+    for (unsigned j = start, end = start + map.value().size(); j < end; j++) {
+      linear = builder.create<arith::MulIOp>(loc, linear, shape[j]);
+    }
+    // Start collapse.
+    Value val;
+    if (!isExpand)
+      val = srcIndices[i];
+    // Iterate over dimension slice.
+    for (unsigned j = start, end = start + map.value().size(); j < end; j++) {
+      linear = builder.create<arith::DivUIOp>(loc, linear, shape[j]);
+      if (isExpand) {
+        Value old = srcIndices[j];
+        Value mul = builder.create<arith::MulIOp>(loc, old, linear);
+        val = val ? builder.create<arith::AddIOp>(loc, val, mul) : mul;
+      } else {
+        Value old = val;
+        val = builder.create<arith::DivUIOp>(loc, val, linear);
+        assert(dstIndices.size() == j);
+        dstIndices.push_back(val);
+        val = builder.create<arith::RemUIOp>(loc, old, linear);
+      }
+    }
+    // Finalize expansion.
+    if (isExpand) {
+      assert(dstIndices.size() == i);
+      dstIndices.push_back(val);
+    }
+    start += map.value().size();
+    i++;
+  }
+  assert(dstIndices.size() == dstRank);
+}

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

 /// TODO: provide as general utility to MLIR at large?
 static void translateIndices(Location loc, ConversionPatternRewriter &rewriter,
                              ArrayRef<ReassociationIndices> reassociation,
                              TensorType dstTp, TensorType srcTp, Value dstIdx,
                              Value srcIdx, ArrayRef<Value> dstShape,
                              ArrayRef<Value> srcShape) {
   unsigned dstRank = dstTp.getRank();
   unsigned srcRank = srcTp.getRank();
-  unsigned start = 0;
-  unsigned i = 0;
-  bool isExpand = srcRank > dstRank;
-  ArrayRef<Value> shape = isExpand ? srcShape : dstShape;
-  // Iterate over reassociation map.
-  for (const auto &map : llvm::enumerate(reassociation)) {
-    // Prepare strides information in dimension slice.
-    Value linear = constantIndex(rewriter, loc, 1);
-    for (unsigned j = start, end = start + map.value().size(); j < end; j++) {
-      linear = rewriter.create<arith::MulIOp>(loc, linear, shape[j]);
-    }
-    // Start collapse.
-    Value idx = constantIndex(rewriter, loc, i++);
-    Value val;
-    if (!isExpand)
-      val = rewriter.create<memref::LoadOp>(loc, srcIdx, idx);
-    // Iterate over dimension slice.
-    for (unsigned j = start, end = start + map.value().size(); j < end; j++) {
-      linear = rewriter.create<arith::DivUIOp>(loc, linear, shape[j]);
-      Value jdx = constantIndex(rewriter, loc, j);
-      if (isExpand) {
-        Value old = rewriter.create<memref::LoadOp>(loc, srcIdx, jdx);
-        Value mul = rewriter.create<arith::MulIOp>(loc, old, linear);
-        val = val ? rewriter.create<arith::AddIOp>(loc, val, mul) : mul;
-      } else {
-        Value old = val;
-        val = rewriter.create<arith::DivUIOp>(loc, val, linear);
-        rewriter.create<memref::StoreOp>(loc, val, dstIdx, jdx);
-        val = rewriter.create<arith::RemUIOp>(loc, old, linear);
-      }
-    }
-    // Finalize expansion.
-    if (isExpand)
-      rewriter.create<memref::StoreOp>(loc, val, dstIdx, idx);
-    start += map.value().size();
+
+  SmallVector<Value, 4> srcIndices;
+  for (unsigned i = 0; i < srcRank; i++) {
+    Value idx = rewriter.create<memref::LoadOp>(
+        loc, srcIdx, constantIndex(rewriter, loc, i));
+    srcIndices.push_back(idx);
   }
-  // Sanity.
-  assert((isExpand && i == dstRank) || (!isExpand && i == srcRank));
+
+  SmallVector<Value, 4> dstIndices;
+  translateIndicesArray(rewriter, loc, reassociation, srcIndices, srcShape,
+                        dstShape, dstIndices);
+
+  for (unsigned i = 0; i < dstRank; i++)
+    rewriter.create<memref::StoreOp>(loc, dstIndices[i], dstIdx,
+                                     constantIndex(rewriter, loc, i));
 }
 
 /// Helper method to compute the shape of destination tensor of a reshape
 /// operator. This is only used when operands have dynamic shape. The shape of
 /// the destination is stored into dstShape.
 void genReshapeDstShape(Location loc, ConversionPatternRewriter &rewriter,
                         SmallVector<Value, 4> &dstShape,
                         ArrayRef<Value> srcShape,

mlir/test/Dialect/SparseTensor/sparse_reshape.mlir

 // CHECK-CONV-DAG:  call @newSparseTensor
 // CHECK-CONV-DAG:  call @newSparseTensor
 // CHECK-CONV:      scf.while : () -> () {
 // CHECK-CONV:        call @getNextF64
 // CHECK-CONV:        scf.condition
 // CHECK-CONV:      } do {
 // CHECK-CONV:        %[[X:.*]] = memref.load %{{.*}}[%[[C0]]] : memref<1xindex>
 // CHECK-CONV:        %[[D:.*]] = arith.divui %[[X]], %[[C10]] : index
-// CHECK-CONV:        memref.store %[[D]], %{{.*}}[%[[C0]]] : memref<2xindex>
 // CHECK-CONV:        %[[R:.*]] = arith.remui %[[X]], %[[C10]] : index
+// CHECK-CONV:        memref.store %[[D]], %{{.*}}[%[[C0]]] : memref<2xindex>
 // CHECK-CONV:        memref.store %[[R]], %{{.*}}[%[[C1]]] : memref<2xindex>
 // CHECK-CONV:        call @addEltF64
 // CHECK-CONV:        scf.yield
 // CHECK-CONV:      }
 // CHECK-CONV:      %[[N:.*]] = call @newSparseTensor
 // CHECK-CONV:      call @delSparseTensorCOOF64
 // CHECK-CONV:      call @delSparseTensorCOOF64
 // CHECK-CONV:      return %[[N]] : !llvm.ptr<i8>

 // CHECK-CONV-DAG:  %[[C10:.*]] = arith.constant 10 : index
 // CHECK-CONV-DAG:  call @newSparseTensor
 // CHECK-CONV-DAG:  call @newSparseTensor
 // CHECK-CONV:      scf.while : () -> () {
 // CHECK-CONV:        call @getNextF64
 // CHECK-CONV:        scf.condition
 // CHECK-CONV:      } do {
 // CHECK-CONV:        %[[X:.*]] = memref.load %{{.*}}[%[[C0]]] : memref<2xindex>
-// CHECK-CONV:        %[[M:.*]] = arith.muli %[[X]], %[[C10]] : index
 // CHECK-CONV:        %[[Y:.*]] = memref.load %{{.*}}[%[[C1]]] : memref<2xindex>
+// CHECK-CONV:        %[[M:.*]] = arith.muli %[[X]], %[[C10]] : index
 // CHECK-CONV:        %[[A:.*]] = arith.addi %[[M]], %[[Y]] : index
 // CHECK-CONV:        memref.store %[[A]], %{{.*}}[%[[C0]]] : memref<1xindex>
 // CHECK-CONV:        call @addEltF64
 // CHECK-CONV:        scf.yield
 // CHECK-CONV:      }
 // CHECK-CONV:      %[[N:.*]] = call @newSparseTensor
 // CHECK-CONV:      call @delSparseTensorCOOF64
 // CHECK-CONV:      call @delSparseTensorCOOF64

 // CHECK-CONV-DAG:  %[[C10:.*]] = arith.constant 10 : index
 // CHECK-CONV-DAG:  %[[D1:.*]] = arith.divui %{{.*}}, %[[C10]] : index
 // CHECK-CONV-DAG:  call @newSparseTensor
 // CHECK-CONV-DAG:  call @newSparseTensor
 // CHECK-CONV:      scf.while : () -> () {
 // CHECK-CONV:        call @getNextF64
 // CHECK-CONV:        scf.condition
 // CHECK-CONV:      } do {
-// CHECK-CONV:        %[[M:.*]] = arith.muli %[[D1]], %[[C10]] : index
 // CHECK-CONV:        %[[L:.*]] = memref.load %{{.*}}[%[[C0]]] : memref<1xindex>
+// CHECK-CONV:        %[[M:.*]] = arith.muli %[[D1]], %[[C10]] : index
 // CHECK-CONV:        %[[D2:.*]] = arith.divui %[[M]], %[[D1]] : index
 // CHECK-CONV:        %[[D3:.*]] = arith.divui %[[L]], %[[D2]] : index
-// CHECK-CONV:        memref.store %[[D3]], %{{.*}}[%[[C0]]] : memref<2xindex>
 // CHECK-CONV:        %[[R:.*]] = arith.remui %[[L]], %[[D2]] : index
 // CHECK-CONV:        %[[D4:.*]] = arith.divui %[[D2]], %[[C10]] : index
 // CHECK-CONV:        %[[D5:.*]] = arith.divui %[[R]], %[[D4]] : index
+// CHECK-CONV:        memref.store %[[D3]], %{{.*}}[%[[C0]]] : memref<2xindex>
 // CHECK-CONV:        memref.store %[[D5]], %{{.*}}[%[[C1]]] : memref<2xindex>
 // CHECK-CONV:        call @addEltF64
 // CHECK-CONV:        scf.yield
 // CHECK-CONV:      }
 // CHECK-CONV:      %[[N:.*]] = call @newSparseTensor
 // CHECK-CONV:      call @delSparseTensorCOOF64
 // CHECK-CONV:      call @delSparseTensorCOOF64
 // CHECK-CONV:      return %[[N]] : !llvm.ptr<i8>

 // CHECK-CONV-DAG:  %[[C10:.*]] = arith.constant 10 : index
 // CHECK-CONV-DAG:  %[[M1:.*]] = arith.muli %{{.*}}, %[[C10]] : index
 // CHECK-CONV-DAG:  call @newSparseTensor
 // CHECK-CONV-DAG:  call @newSparseTensor
 // CHECK-CONV:      scf.while : () -> () {
 // CHECK-CONV:        call @getNextF64
 // CHECK-CONV:        scf.condition
 // CHECK-CONV:      } do {
-// CHECK-CONV:        %[[D1:.*]] = arith.divui %[[M1]], %[[C10]] : index
 // CHECK-CONV:        %[[X:.*]] = memref.load %{{.*}}[%[[C0]]] : memref<2xindex>
+// CHECK-CONV:        %[[Y:.*]] = memref.load %{{.*}}[%[[C1]]] : memref<2xindex>
+// CHECK-CONV:        %[[D1:.*]] = arith.divui %[[M1]], %[[C10]] : index
 // CHECK-CONV:        %[[M2:.*]] = arith.muli %[[X]], %[[D1]] : index
 // CHECK-CONV:        %[[D2:.*]] = arith.divui %[[D1]], %{{.*}} : index
-// CHECK-CONV:        %[[Y:.*]] = memref.load %{{.*}}[%[[C1]]] : memref<2xindex>
 // CHECK-CONV:        %[[M3:.*]] = arith.muli %[[Y]], %[[D2]] : index
 // CHECK-CONV:        %[[A:.*]] = arith.addi %[[M2]], %[[M3]] : index
 // CHECK-CONV:        memref.store %[[A]], %{{.*}}[%[[C0]]] : memref<1xindex>
 // CHECK-CONV:        call @addEltF64
 // CHECK-CONV:        scf.yield
 // CHECK-CONV:      }
 // CHECK-CONV:      %[[N:.*]] = call @newSparseTensor
 // CHECK-CONV:      call @delSparseTensorCOOF64
 // CHECK-CONV:      call @delSparseTensorCOOF64
 // CHECK-CONV:      return %[[N]] : !llvm.ptr<i8>
 //
 func.func @dynamic_sparse_collapse(%arg0: tensor<10x?xf64, #SparseMatrix>) -> tensor<?xf64, #SparseVector> {
   %0 = tensor.collapse_shape %arg0 [[0, 1]] :
     tensor<10x?xf64, #SparseMatrix> into tensor<?xf64, #SparseVector>
   return %0 : tensor<?xf64, #SparseVector>
 }