[Triton Tool] Enable the triton-tensor-layout tool for 3rd party Trit…

…on GPU Dialect (#5444)

Remove the protection of only the `ttg` dialect being allowed in
triton-tensor-layout tool safely.

Change the helper function to check whether the layout attribute
implements the `toLinearLayout` interface.

This PR does not include new test because the in-tree LIT test covers
the changes for in-tree layout attribute of Triton GPU dialect.

bin/triton-tensor-layout.cpp

@@ -80,10 +80,11 @@ static cl::opt<std::string> TensorStr(
 //===--------------------------------------------------------------------===//
 
 LogicalResult layoutPrint(RankedTensorType tensorType, raw_ostream &os) {
-  StringRef dialectName = tensorType.getEncoding().getDialect().getNamespace();
-
-  // Dispatch to the corresponding dialect helper function to print the layout.
-  if (dialectName == "ttg") {
+  // DistributedEncodingTrait and SharedEncodingAttr implements the
+  // toLinearLayout interface.
+  mlir::Attribute layout = tensorType.getEncoding();
+  if (isa<mlir::triton::gpu::DistributedEncodingTrait,
+          mlir::triton::gpu::SharedEncodingAttr>(layout)) {
     os << triton::gpu::getLayoutStr(tensorType, UseHWPointOfView);
     return success();
   }