[mlir][sparse] add narrower choices for pointers/indices

[aartbik]

Summary

Use cases with 16- or even 8-bit pointer/index structures have been identified.

[penpornk]

Thank you for adding this, Aart! :)

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

517

Is the type signed by default when isSigned is not specified? I think we want unsigned int8 too.

[aartbik]

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

517

These are "signless" (MLIR knows signed, unsigned, and signless). I indeed don't like that the lowering of the cast from signless to index uses a sign extension (llvm.sext), but explicitly unsigned integers are cast incompatible with index for some reason.

[aartbik]

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

517

Signless really leaves it to the operator to decide what interpretation is the best. So I think I should use "zexti" instead of the index_cast, and we are good!

I will do that in a follow-up CL though....

mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h

 /// a co-iterating while-loop).
 enum class SparseVectorizationStrategy {
   kNone,
   kDenseInnerLoop,
   kAnyStorageInnerLoop
 };
 
 /// Defines a type for "pointer" and "index" storage in the sparse storage
-/// scheme, with a choice between the native platform-dependent index width,
-/// 64-bit integers, or 32-bit integers. A narrow width obviously reduces
+/// scheme, with a choice between the native platform-dependent index width
+/// or any of 64-/32-/16-/8-bit integers. A narrow width obviously reduces
 /// the memory footprint of the sparse storage scheme, but the width should
 /// suffice to define the total required range (viz. the maximum number of
 /// stored entries per indirection level for the "pointers" and the maximum
 /// value of each tensor index over all dimensions for the "indices").
-enum class SparseIntType { kNative, kI64, kI32 };
+enum class SparseIntType { kNative, kI64, kI32, kI16, kI8 };
 
 /// Sparsification options.
 struct SparsificationOptions {
   SparsificationOptions(SparseParallelizationStrategy p,
                         SparseVectorizationStrategy v, unsigned vl,
                         SparseIntType pt, SparseIntType it)
       : parallelizationStrategy(p), vectorizationStrategy(v), vectorLength(vl),
         ptrType(pt), indType(it) {

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

 static Type genIntType(PatternRewriter &rewriter, linalg::SparseIntType tp) {
   switch (tp) {
   case linalg::SparseIntType::kNative:
     return rewriter.getIndexType();
   case linalg::SparseIntType::kI64:
     return rewriter.getIntegerType(64);
   case linalg::SparseIntType::kI32:
     return rewriter.getIntegerType(32);
+  case linalg::SparseIntType::kI16:
+    return rewriter.getIntegerType(16);
+  case linalg::SparseIntType::kI8:
+    return rewriter.getIntegerType(8);
   }
 }
 
 /// 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) {

mlir/test/Dialect/Linalg/sparse_storage.mlir

 // RUN: mlir-opt %s -test-sparsification="ptr-type=1 ind-type=1" | \
 // RUN:   FileCheck %s --check-prefix=CHECK-TYPE0
 // RUN: mlir-opt %s -test-sparsification="ptr-type=1 ind-type=2" | \
 // RUN:   FileCheck %s --check-prefix=CHECK-TYPE1
 // RUN: mlir-opt %s -test-sparsification="ptr-type=2 ind-type=1" | \
 // RUN:   FileCheck %s --check-prefix=CHECK-TYPE2
 // RUN: mlir-opt %s -test-sparsification="ptr-type=2 ind-type=2" | \
 // RUN:   FileCheck %s --check-prefix=CHECK-TYPE3
+// RUN: mlir-opt %s -test-sparsification="ptr-type=3 ind-type=3" | \
+// RUN:   FileCheck %s --check-prefix=CHECK-TYPE4
+// RUN: mlir-opt %s -test-sparsification="ptr-type=4 ind-type=4" | \
+// RUN:   FileCheck %s --check-prefix=CHECK-TYPE5
 
 #trait_mul_1d = {
   indexing_maps = [
     affine_map<(i) -> (i)>,  // a
     affine_map<(i) -> (i)>,  // b
     affine_map<(i) -> (i)>   // x (out)
   ],
   sparse = [

 // CHECK-TYPE3:   %[[IND0:.*]] = load %{{.*}}[%[[I]]] : memref<?xi32>
 // CHECK-TYPE3:   %[[INDC:.*]] = index_cast %[[IND0]] : i32 to index
 // CHECK-TYPE3:   %[[VAL0:.*]] = load %{{.*}}[%[[I]]] : memref<?xf64>
 // CHECK-TYPE3:   %[[VAL1:.*]] = load %{{.*}}[%[[INDC]]] : memref<32xf64>
 // CHECK-TYPE3:   %[[MUL:.*]] = mulf %[[VAL0]], %[[VAL1]] : f64
 // CHECK-TYPE3:   store %[[MUL]], %{{.*}}[%[[INDC]]] : memref<32xf64>
 // CHECK-TYPE3: }
 
+// CHECK-TYPE4-LABEL: func @mul_dd(
+// CHECK-TYPE4: %[[C0:.*]] = constant 0 : index
+// CHECK-TYPE4: %[[C1:.*]] = constant 1 : index
+// CHECK-TYPE4: %[[P0:.*]] = load %{{.*}}[%[[C0]]] : memref<?xi16>
+// CHECK-TYPE4: %[[B0:.*]] = index_cast %[[P0]] : i16 to index
+// CHECK-TYPE4: %[[P1:.*]] = load %{{.*}}[%[[C1]]] : memref<?xi16>
+// CHECK-TYPE4: %[[B1:.*]] = index_cast %[[P1]] : i16 to index
+// CHECK-TYPE4: scf.for %[[I:.*]] = %[[B0]] to %[[B1]] step %[[C1]] {
+// CHECK-TYPE4:   %[[IND0:.*]] = load %{{.*}}[%[[I]]] : memref<?xi16>
+// CHECK-TYPE4:   %[[INDC:.*]] = index_cast %[[IND0]] : i16 to index
+// CHECK-TYPE4:   %[[VAL0:.*]] = load %{{.*}}[%[[I]]] : memref<?xf64>
+// CHECK-TYPE4:   %[[VAL1:.*]] = load %{{.*}}[%[[INDC]]] : memref<32xf64>
+// CHECK-TYPE4:   %[[MUL:.*]] = mulf %[[VAL0]], %[[VAL1]] : f64
+// CHECK-TYPE4:   store %[[MUL]], %{{.*}}[%[[INDC]]] : memref<32xf64>
+// CHECK-TYPE4: }
+
+// CHECK-TYPE5-LABEL: func @mul_dd(
+// CHECK-TYPE5: %[[C0:.*]] = constant 0 : index
+// CHECK-TYPE5: %[[C1:.*]] = constant 1 : index
+// CHECK-TYPE5: %[[P0:.*]] = load %{{.*}}[%[[C0]]] : memref<?xi8>
+// CHECK-TYPE5: %[[B0:.*]] = index_cast %[[P0]] : i8 to index
+// CHECK-TYPE5: %[[P1:.*]] = load %{{.*}}[%[[C1]]] : memref<?xi8>
+// CHECK-TYPE5: %[[B1:.*]] = index_cast %[[P1]] : i8 to index
+// CHECK-TYPE5: scf.for %[[I:.*]] = %[[B0]] to %[[B1]] step %[[C1]] {
+// CHECK-TYPE5:   %[[IND0:.*]] = load %{{.*}}[%[[I]]] : memref<?xi8>
+// CHECK-TYPE5:   %[[INDC:.*]] = index_cast %[[IND0]] : i8 to index
+// CHECK-TYPE5:   %[[VAL0:.*]] = load %{{.*}}[%[[I]]] : memref<?xf64>
+// CHECK-TYPE5:   %[[VAL1:.*]] = load %{{.*}}[%[[INDC]]] : memref<32xf64>
+// CHECK-TYPE5:   %[[MUL:.*]] = mulf %[[VAL0]], %[[VAL1]] : f64
+// CHECK-TYPE5:   store %[[MUL]], %{{.*}}[%[[INDC]]] : memref<32xf64>
+// CHECK-TYPE5: }
+
 func @mul_dd(%arga: tensor<32xf64>, %argb: tensor<32xf64>) -> tensor<32xf64> {
   %0 = linalg.generic #trait_mul_1d
      ins(%arga, %argb: tensor<32xf64>, tensor<32xf64>)
     outs(%arga : tensor<32xf64>) {
       ^bb(%a: f64, %b: f64, %s: f64):
         %0 = mulf %a, %b  : f64
         linalg.yield %0 : f64
   } -> tensor<32xf64>
   return %0 : tensor<32xf64>
 }
 

mlir/test/lib/Transforms/TestSparsification.cpp

struct TestSparsification
   linalg::SparseIntType typeOption(int32_t option) {
     switch (option) {
     default:
       return linalg::SparseIntType::kNative;
     case 1:
       return linalg::SparseIntType::kI64;
     case 2:
       return linalg::SparseIntType::kI32;
+    case 3:
+      return linalg::SparseIntType::kI16;
+    case 4:
+      return linalg::SparseIntType::kI8;
     }
   }
 
   /// Runs the test on a function.
   void runOnFunction() override {
     auto *ctx = &getContext();
     OwningRewritePatternList patterns;
     // Translate strategy flags to strategy options.