RFC: Generating code to filter out explicit zeros in results

[stephenchouca]

Summary

This RFC proposes adding a new padded property to mode formats, which describes whether or not explicit zeros are stored. The lowering machinery can use this property to determine whether/how to generate code that filters out explicit zeros in the result.

Motivation

Currently, an expression such as a(i) = b(i), where a is compressed and b is dense, is always lowered to the following code:

// Algorithm 1
int32_t pa1 = 0;
for (int32_t i = 0; i < b1_dimension; i++) {
    a_vals[pa1] = b_vals[i];
    pa1++;
}

which stores all components of b into a, including all the zero components. This is often not desirable, since the whole point of having compressed formats in the first place is to avoid needing to store zeros. Instead, we might want to generate the following code:

// Algorithm 2
int32_t pa1 = 0;
for (int32_t i = 0; i < b1_dimension; i++) {
    if (b_vals[i] != 0.0) {
        a_vals[pa1] = b_vals[i];
        pa1++;
    }
}

which only stores non-zero components of b into a. This is clearly beneficial if most components of b are zeros. However, having to filter out zeros incurs additional performance overhead, so if we know that b contains only (or even mostly) nonzeros then Algorithm 1 would actually be preferable. Thus, we want TACO to be able to generate both Algorithms 1 and 2.

Proposed Changes

I propose adding a new padded property to mode formats. A padded mode format may store explicit zero values (or subtensors), whereas an unpadded mode format only stores nonzero values (or subtensors). The lowering machinery can then use this property to determine whether (and how) to generate code that filters out explicit zeros in the result. For instance:

• If a from the earlier example is stored in a padded version of the compressed mode format, then the result is allowed to store explicit zeros.
  Thus, the lowering machinery can emit Algorithm 1 for the vector assignment, since explicit zeros do not need to be filtered out.
• If a is stored in an unpadded version of the compressed mode format and b is stored in a padded version of the dense mode format, then the input may contain explicit zeros but the result must not store explicit zeros.
  Thus, the lowering machinery must emit Algorithm 2 for the vector assignment in order to filter out explicit zeros.
• If b is also stored in an unpadded version of the dense mode format though, then the lowering machinery can infer that the result will not contain zeros.
  Thus, the lowering machinery can again emit Algorithm 1 for the vector assignment, since there will not be explicit zeros that need to be filtered out.

For matrices and higher-order tensors, arbitrary combinations of padded and unpadded mode formats may be used to store tensor dimensions. CSR without explicit zeros, for instance, can be expressed as (dense, compressed) where the dense row dimension is padded (to indicate that empty rows may still be stored in the pos array) and the compressed column dimension is unpadded. By contrast, using padded compressed to store the column dimension describes CSR with explicit zeros, while using unpadded dense to store the row dimension describes a CSR matrix with no empty row.

[dcbdan]

Being able to call a(i) = b(i) where a is compressed and b is dense, and generalizations thereof, would be an important addition. It would also be important to do something akin to

  a(i) = mask(b(i) <= 0.0, b(i))     // relu
  a(i) = mask(abs(b(i)) < eps, b(i)) // threshold

where mask takes in a boolean expression where if true, the output is zero.

Questions:

• How would the padded mode attribute work for thresholding operators?
• It isn't clear to me that it would be guaranteed that a dense unpadded mode contains no zeros.
  Is it necessary to add unpadded/padded attribute to dense modes?
• Is more complicated rhs a concern? As in a(i) = b(i) + c(i) or a(i) = B(i,j)*c(j)?

[dcbdan]

What is the status of these potential changes?

[stephenchouca]

@dcbdan Thanks for your comments, and sorry for the late reply. With regards to your questions, I believe it should be possible to support filtering out explicit zeros in the result even with more complicated RHS expressions. For some expressions, like sparse vector addition, it should be as straightforward as just inserting guards into the generated code to store to the result only if the sum if non-zero:

 while (pb1 < pb1_end && pc1 < pc1_end) {
    int32_t ib = b1_crd[pb1];
    int32_t ic = c1_crd[pc1];
    int32_t i = TACO_MIN(ib,ic);
    if (ib == i && ic == i) {
      double sum = b_vals[pb1] + c_vals[pc1];
      if (sum != 0.0) {
        a_vals[pa1]
        pa1++;
      }
    }
    else if (ib == i) {
      if (b_vals[pb1] != 0.0) {
        a_vals[pa1] = b_vals[pb1];
        pa1++;
      }
    }
    else {
      if (c_vals[pc1] != 0.0) {
        a_vals[pa1] = c_vals[pc1];
        pa1++;
      }
    }
    pb1 += (int32_t)(ib == i);
    pc1 += (int32_t)(ic == i);
  }
  while (pb1 < pb1_end) {
    if (b_vals[pb1] != 0.0) {
      a_vals[pa1] = b_vals[pb1];
      pa1++;
    }
    pb1++;
  }
  while (pc1 < pc1_end) {
    if (c_vals[pc1] != 0.0) {
      a_vals[pa1] = c_vals[pc1];
      pa1++;
    }
    pc1++;
  }

For other expressions that scatter into the result, such as sparse matrix-sparse vector multiplication with CSC matrices, it might also require generating workspaces (#184) for storing temporary results with zeros, which could then get filtered out when the workspace is copied over to the final result. This would also work for thresholding operations if we express them as something like a(i) = (abs(b(i)) < eps) * b(i). As for allowing unpadded dense modes, while it wouldn't be possible to guarantee that results do not contain explicit zeros, unpadded dense modes could still be used by users to optionally specify that inputs to an expression do not explicit zeros (see the examples in my original proposal). Without such information, TACO would be unable to generate optimized code that omits unnecessary checks if explicit zeros need to be filtered out but the RHS can be proven to not ever be zero.

I plan to start looking into implementing this proposal later this week. Fully implementing this might take several weeks, though depending on whether I run into unexpected issues it might take longer (or shorter).