[Mlir-commits] [mlir] [mlir] Add unit test for RankedTensorType wrapper class example. (PR #99789)

Jacques Pienaar llvmlistbot at llvm.org
Mon Jul 22 12:49:27 PDT 2024


================
@@ -226,4 +227,63 @@ TEST(ShapedTypeTest, RankedTensorTypeBuilder) {
   }
 }
 
+/// Simple wrapper class to enable "isa querying" and simple accessing of
+/// encoding.
+class TensorWithString : public RankedTensorType {
+public:
+  using RankedTensorType::RankedTensorType;
+
+  static TensorWithString get(ArrayRef<int64_t> shape, Type elementType,
+                              StringRef name) {
+    return mlir::cast<TensorWithString>(RankedTensorType::get(
+        shape, elementType, StringAttr::get(elementType.getContext(), name)));
+  }
+
+  StringRef getName() const {
+    if (Attribute enc = getEncoding())
+      return mlir::cast<StringAttr>(enc).getValue();
+    return {};
+  }
+
+  static bool classof(Type type);
+};
+
+bool TensorWithString::classof(Type type) {
+  if (auto rt = mlir::dyn_cast_or_null<RankedTensorType>(type))
+    return mlir::isa_and_present<StringAttr>(rt.getEncoding());
+  return false;
+}
----------------
jpienaar wrote:

Made inline for simplicity (no good reason here).

https://github.com/llvm/llvm-project/pull/99789


More information about the Mlir-commits mailing list