[mlir][sparse] Improving SparseTensorDimSliceAttr methods

[wrengr]

Summary

This patch makes the following changes to SparseTensorDimSliceAttr methods:

• Mark isDynamic constexpr.
• Add new helpers getStatic and getStaticString to avoid repetition.
• Moved the definitions for getStatic{Offset,Stride,Size} and isCompletelyDynamic out of the class declaration; because there's no benefit to inlining them.
• Changed parse to use kDynamic rather than literals.
• Changed verify to use the isDynamic helper.

[Peiming]

mlir/lib/Dialect/SparseTensor/IR/SparseTensorDialect.cpp

221

Why call getStaticString when the v can be dynamic?

[wrengr]

mlir/lib/Dialect/SparseTensor/IR/SparseTensorDialect.cpp

221

Perhaps it's not the best name, but here's the reasoning that lead to it: When factoring out the code duplication part of getStatic{Offset,Stride,Size}, it makes sense to call the new method "getStatic". In English, the duplicate code in the print method was essentially saying "if getV() is static, then to_string it, else return '?'". In pseudocode we can rephrase that to "getString(getStatic(getV()))". But then rather than defining "getString" it'd make more sense to define the composite getString*getStatic. And that composite is just like getStatic except for returning a string, hence getStaticString.

The big naming problem here is that we don't have a name/class that comprises all three of {Offset,Stride,Size}. If we did, then we could name it "DSA::fooToString(Foo)", "DSA::getFooString(Foo)", "Foo::getString()", or similar.

I'm open to suggestions if you have a better name in mind

[Peiming]

mlir/lib/Dialect/SparseTensor/IR/SparseTensorDialect.cpp

221

I also don’t think getStaticString should be a member function either. (I can not think of other place where this function is needed outside the op::print). It might better be a file local utils.

BTW, I also find that ShapedType has both kDynamicSize and KDynamicOffsetOrStride, is it better to reuse those here?

[wrengr]

mlir/lib/Dialect/SparseTensor/IR/SparseTensorDialect.cpp

221

Yeah, ideally I feel like all four of {kDynamic,isDynamic,getStatic,getStaticString} should really just be file-local rather than being part of the class, since none of them are needed outside of the {parse,verify,print} methods. The only reason I didn't do things that way was because I couldn't think of decent names for them, so keeping them under the SparseTensorDimSliceAttr namespace avoids the naming problem.

Again, I'm open to other suggestions. As the description says, the main goal of the CL is just to: (1) move defns out of the class decl; (2) reduce repetition; (3) have {parse,verify} use the helpers, not just for reducing repetition but also for ensuring that they remain in sync (e.g., if we decided to change kDynamic to being INT_MIN, or decided to change these parameters into newtypes or attributes, etc). So, I'm not wedded to the names or where the helpers live, just so long as they satisfy those goals.

221

BTW, I also find that ShapedType has both kDynamicSize and KDynamicOffsetOrStride, is it better to reuse those here?

As for whether to reuse ShapedType::kDynamic{Size,OffsetOrStride}: I think that's a good idea, since those're the abstractions used elsewhere in MLIR for the same task, therefore using them ourselves would help reduce conceptual complexity (and maybe even code duplication).

(Fwiw, I do have a CL in the wings for defining a new Size type that captures a bunch of the special logic associated with sizes. Whenever I get back to finishing it, it's the sort of thing I think should be upstreamed and reused by other dialects. I originally started working on it in order to simplify the code for the dim<->lvl stuff; but since it's a fair deal of code, I eventually figured out a way to remove it from the rest of the dim<->lvl stuff without introducing too much code duplication. I'm not suggesting the SparseTensorDimSliceAttr changes should wait until the Size code is in; just giving a random heads-up is all.)

[wrengr]

rebase

[wrengr]

rebasing over D151505, D151513, D151518, D151522

mlir/include/mlir/Dialect/SparseTensor/IR/SparseTensorAttrDefs.td

let parameters = (
     "int64_t" : $offset,
     "int64_t" : $size,
     "int64_t" : $stride
   );
 
   let extraClassDeclaration = [{
     /// Special value for dynamic offset/size/stride.
     static constexpr int64_t kDynamic = -1;
-
-    static bool isDynamic(int64_t v) {
-      return v == kDynamic;
-    }
-
-    std::optional<uint64_t> getStaticOffset() const {
-      if (isDynamic(getOffset()))
-        return std::nullopt;
-      return static_cast<uint64_t>(getOffset());
-    };
-
-    std::optional<uint64_t> getStaticStride() const {
-      if (isDynamic(getStride()))
-        return std::nullopt;
-      return static_cast<uint64_t>(getStride());
-    }
-
-    std::optional<uint64_t> getStaticSize() const {
-      if (isDynamic(getSize()))
-        return std::nullopt;
-      return static_cast<uint64_t>(getSize());
-    }
-
-    bool isCompletelyDynamic() const {
-      return isDynamic(getOffset()) && isDynamic(getStride()) && isDynamic(getSize());
-    };
+    static constexpr bool isDynamic(int64_t v) { return v == kDynamic; }
+    static std::optional<uint64_t> getStatic(int64_t v);
+    static std::string getStaticString(int64_t v);
+
+    std::optional<uint64_t> getStaticOffset() const;
+    std::optional<uint64_t> getStaticStride() const;
+    std::optional<uint64_t> getStaticSize() const;
+    bool isCompletelyDynamic() const;
   }];
 
   let genVerifyDecl = 1;
   let hasCustomAssemblyFormat = 1;
 }
 
 //===----------------------------------------------------------------------===//
 // Sparse Tensor Type Encoding Attribute.

mlir/lib/Dialect/SparseTensor/IR/SparseTensorDialect.cpp

 
 /// Gets the dimension-rank of the type of some `T`.  (In particular
 /// this is only used for `Value` and `TypedValue<RankedTensorType>`.)
 template <typename T>
 static inline Dimension getDimRank(T t) {
   return getRankedTensorType(t).getRank();
 }
 
+static constexpr bool acceptBitWidth(unsigned bitWidth) {
+  switch (bitWidth) {
+  case 0:
+  case 8:
+  case 16:
+  case 32:
+  case 64:
+    return true;
+  default:
+    return false;
+  }
+}
+
 //===----------------------------------------------------------------------===//
 // StorageLayout
 //===----------------------------------------------------------------------===//
 
 static constexpr Level kInvalidLevel = -1u;
 static constexpr Level kInvalidFieldIndex = -1u;
 static constexpr FieldIndex kDataFieldStartingIdx = 0;
 
StorageLayout::getFieldIndexAndStride(SparseTensorFieldKind kind,
   assert(fieldIdx != kInvalidFieldIndex);
   return std::pair<FieldIndex, unsigned>(fieldIdx, stride);
 }
 
 //===----------------------------------------------------------------------===//
 // TensorDialect Attribute Methods.
 //===----------------------------------------------------------------------===//
 
-static bool acceptBitWidth(unsigned bitWidth) {
-  switch (bitWidth) {
-  case 0:
-  case 8:
-  case 16:
-  case 32:
-  case 64:
-    return true;
-  default:
-    return false;
-  }
+std::optional<uint64_t> SparseTensorDimSliceAttr::getStatic(int64_t v) {
+  return isDynamic(v) ? std::nullopt
+                      : std::make_optional(static_cast<uint64_t>(v));
+}
+
+std::optional<uint64_t> SparseTensorDimSliceAttr::getStaticOffset() const {
+  return getStatic(getOffset());
+}
+
+std::optional<uint64_t> SparseTensorDimSliceAttr::getStaticStride() const {
+  return getStatic(getStride());
+}
+
+std::optional<uint64_t> SparseTensorDimSliceAttr::getStaticSize() const {
+  return getStatic(getSize());
+}
+
+bool SparseTensorDimSliceAttr::isCompletelyDynamic() const {
+  return isDynamic(getOffset()) && isDynamic(getStride()) &&
+         isDynamic(getSize());
+
+std::string SparseTensorDimSliceAttr::getStaticString(int64_t v) {
+  return isDynamic(v) ? "?" : std::to_string(v);
 }
 
 void SparseTensorDimSliceAttr::print(AsmPrinter &printer) const {
   printer << "(";
-  printer << (getStaticOffset() ? std::to_string(*getStaticOffset()) : "?");
+  printer << getStaticString(getOffset());
   printer << ", ";
-  printer << (getStaticSize() ? std::to_string(*getStaticSize()) : "?");
+  printer << getStaticString(getSize());
   printer << ", ";
-  printer << (getStaticStride() ? std::to_string(*getStaticStride()) : "?");
+  printer << getStaticString(getStride());
   printer << ")";
 }
 
 static ParseResult parseOptionalStaticSlice(int64_t &result,
                                             AsmParser &parser) {
   auto parseResult = parser.parseOptionalInteger(result);
   if (parseResult.has_value()) {
     if (parseResult.value().succeeded() && result < 0) {
       parser.emitError(
           parser.getCurrentLocation(),
           "expect positive value or ? for slice offset/size/stride");
       return failure();
     }
     return parseResult.value();
   }
 
   // Else, and '?' which represented dynamic slice
   result = SparseTensorDimSliceAttr::kDynamic;
   return parser.parseQuestion();
 }
 
 Attribute SparseTensorDimSliceAttr::parse(AsmParser &parser, Type type) {
-  int64_t offset = -1, size = -1, stride = -1;
+  int64_t offset = kDynamic, size = kDynamic, stride = kDynamic;
 
   if (failed(parser.parseLParen()) ||
       failed(parseOptionalStaticSlice(offset, parser)) ||
       failed(parser.parseComma()) ||
       failed(parseOptionalStaticSlice(size, parser)) ||
       failed(parser.parseComma()) ||
       failed(parseOptionalStaticSlice(stride, parser)) ||
       failed(parser.parseRParen()))
     return {};
 
   return parser.getChecked<SparseTensorDimSliceAttr>(parser.getContext(),
                                                      offset, size, stride);
 }
 
 LogicalResult
 SparseTensorDimSliceAttr::verify(function_ref<InFlightDiagnostic()> emitError,
                                  int64_t offset, int64_t size, int64_t stride) {
-  if ((offset == SparseTensorDimSliceAttr::kDynamic || offset >= 0) &&
-      (size == SparseTensorDimSliceAttr::kDynamic || size > 0) &&
-      (stride == SparseTensorDimSliceAttr::kDynamic || stride > 0)) {
-    return success();
-  }
-  return emitError()
-         << "expect positive value or ? for slice offset/size/stride";
-}
+  if (!isDynamic(offset) && offset < 0)
+    return emitError() << "expect non-negative value or ? for slice offset";
+  if (!isDynamic(size) && size <= 0)
+    return emitError() << "expect positive value or ? for slice size";
+  if (!isDynamic(stride) && stride <= 0)
+    return emitError() << "expect positive value or ? for slice stride";
 
 Type mlir::sparse_tensor::detail::getIntegerOrIndexType(MLIRContext *ctx,
                                                         unsigned bitwidth) {
   if (bitwidth)
     return IntegerType::get(ctx, bitwidth);
   return IndexType::get(ctx);
 }