removing mlir ptr dialect and ptr types in codes

Author: Realtyxxx

include/triton-shared/Dialect/TPtr/IR/TPtrDialect.h

@@ -4,8 +4,8 @@
 #include "mlir/Bytecode/BytecodeOpInterface.h"
 #include "mlir/Interfaces/SideEffectInterfaces.h" // Required for IR/TPtrOps.h.inc
 
-#include "mlir/Dialect/Ptr/IR/PtrDialect.h" // Required for IR/TPtrOps.h.inc
-#include "mlir/Dialect/Ptr/IR/PtrTypes.h"   // Required for IR/TPtrOps.h.inc
+// #include "mlir/Dialect/Ptr/IR/PtrDialect.h" // Required for IR/TPtrOps.h.inc
+// #include "mlir/Dialect/Ptr/IR/PtrTypes.h"   // Required for IR/TPtrOps.h.inc
 
 //===----------------------------------------------------------------------===//
 // Temporary Pointer Dialect Operations

include/triton-shared/Dialect/TPtr/IR/TPtrDialect.td

@@ -3,11 +3,14 @@
 
 include "mlir/IR/OpBase.td"
 include "mlir/Interfaces/SideEffectInterfaces.td"
-include "mlir/Dialect/Ptr/IR/MemorySpaceInterfaces.td"
-include "mlir/Dialect/Ptr/IR/PtrDialect.td"
+// include "mlir/Dialect/Ptr/IR/MemorySpaceInterfaces.td"
+// include "mlir/Dialect/Ptr/IR/PtrDialect.td"
 include "mlir/IR/AttrTypeBase.td"
 include "mlir/IR/BuiltinAttributeInterfaces.td"
 include "mlir/IR/BuiltinTypeInterfaces.td"
+include "triton/Dialect/Triton/IR/TritonDialect.td"
+include "triton/Dialect/Triton/IR/TritonTypes.td"
+
 
 def TPtr_Dialect : Dialect {
   let name = "tptr";
@@ -24,8 +27,8 @@ def TPtr_Dialect : Dialect {
     void registerTypes();
   }];
 
-  let dependentDialects = [
-    "mlir::ptr::PtrDialect"
+ let dependentDialects = [
+  "mlir::triton::TritonDialect"
   ];
 
   let useDefaultAttributePrinterParser = 1;
@@ -47,10 +50,10 @@ class TPtr_Attr<string name, string _mnemonic,
 // Memory space attr is required for building Ptr ops
 // This acts as default memory space since there is
 // no such default implemented upstream
-def DefaultMemorySpaceAttr
-    : TPtr_Attr<"DefaultMemorySpace", "default_memory_space",
-                     [DeclareAttrInterfaceMethods<MemorySpaceAttrInterface>]> {
-}
+// def DefaultMemorySpaceAttr
+//     : TPtr_Attr<"DefaultMemorySpace", "default_memory_space",
+//                      [DeclareAttrInterfaceMethods<MemorySpaceAttrInterface>]> {
+// }
 
 //
 // Op Base
@@ -72,7 +75,7 @@ def TPTR_IntToPtrOp : TPTR_Op<"inttoptr", [
     ```
   }];
   let arguments = (ins AnySignlessIntegerOrIndex:$arg);
-  let results = (outs Ptr_PtrType:$res);
+  let results = (outs TT_PtrLike:$res);
   let assemblyFormat = "$arg attr-dict `:` type($arg) `to` type($res)";
 }
 
@@ -88,45 +91,61 @@ def TPTR_PtrToIntOp : TPTR_Op<"ptrtoint", [
     %int = ptr.ptrtoint %ptr : !ptr.ptr<1 : i32> to i32
     ```
   }];
-  let arguments = (ins Ptr_PtrType:$arg);
+  let arguments = (ins TT_PtrLike:$arg);
   let results = (outs AnySignlessIntegerOrIndex:$res);
   let assemblyFormat = "$arg attr-dict `:` type($arg) `to` type($res)";
 }
 
-def TPTR_TypeOffsetOp : TPTR_Op<"type_offset", [ConstantLike, Pure]> {
-  let summary = "Creates a type offset constant.";
-  let description = [{
-    The `addr.type_offset` operation produces an int or index-typed SSA value
-    equal to a target-specific constant representing the offset of a single
-    element of the given type. The default return type is `index`.
-    Example:
-
-    ```mlir
-    %0 = addr.type_offset f32
-    %1 = addr.type_offset memref<12 x f64> : i32
-    ```
-  }];
-
-  let arguments = (ins TypeAttr:$baseType);
-  let results = (outs AnySignlessIntegerOrIndex:$result);
-  let builders = [
-    OpBuilder<(ins "TypeAttr":$baseType, CArg<"Type", "nullptr">:$resultTy)>
-  ];
-  let assemblyFormat = [{
-     attr-dict $baseType custom<IntType>(type($result))
-  }];
-  let hasFolder = 1;
-}
+// def TPTR_TypeOffsetOp : TPTR_Op<"type_offset", [ConstantLike, Pure]> {
+//   let summary = "Creates a type offset constant.";
+//   let description = [{
+//     The `addr.type_offset` operation produces an int or index-typed SSA value
+//     equal to a target-specific constant representing the offset of a single
+//     element of the given type. The default return type is `index`.
+//     Example:
+
+//     ```mlir
+//     %0 = addr.type_offset f32
+//     %1 = addr.type_offset memref<12 x f64> : i32
+//     ```
+//   }];
+
+//   let arguments = (ins TypeAttr:$baseType);
+//   let results = (outs AnySignlessIntegerOrIndex:$result);
+//   let builders = [
+//     OpBuilder<(ins "TypeAttr":$baseType, CArg<"Type", "nullptr">:$resultTy), [{
+//       Type resultType = resultTy ? resultTy : $_builder.getIndexType();
+//       $_state.addAttribute("baseType", baseType);
+//       $_state.addTypes(resultType);
+//     }]>,
+//     OpBuilder<(ins "Type":$baseType), [{
+//       TypeAttr baseTypeAttr = TypeAttr::get(baseType);
+//       Type resultType = $_builder.getIndexType();
+//       $_state.addAttribute("baseType", baseTypeAttr);
+//       $_state.addTypes(resultType);
+//     }]>
+//   ];
+//   let assemblyFormat = [{
+//      attr-dict $baseType custom<IntType>(type($result))
+//   }];
+// let extraClassDeclaration = [{
+//     /// Returns the type offset according to `layout`. If `layout` is `nullopt`
+//     /// the nearest layout the op will be used for the computation.
+//     /// ! I copied this from ptr dialect
+//     llvm::TypeSize getTypeSize(std::optional<DataLayout> layout = std::nullopt);
+//   }];
+//   let hasFolder = 1;
+// }
 
 def TPTR_FromMemrefOp : TPTR_Op<"from_memref", [Pure]> {
   let arguments = (ins AnyMemRef:$input);
-  let results = (outs Ptr_PtrType:$result);
+  let results = (outs TT_PtrLike:$result);
   let assemblyFormat = "$input attr-dict `:` type($input) `to` type($result)";
 }
 
 def TPTR_ToMemrefOp : TPTR_Op<"to_memref", [
     Pure  ]> {
-  let arguments = (ins Ptr_PtrType:$arg);
+  let arguments = (ins TT_PtrLike:$arg);
   let results = (outs AnyStaticShapeMemRef:$res);
   let assemblyFormat = "$arg attr-dict `:` type($arg) `to` type($res)";
 }
@@ -143,8 +162,8 @@ def TPTR_PtrAddOp : TPTR_Op<"ptradd", [Pure, AllTypesMatch<["base", "result"]>]>
     ```
   }];
 
-  let arguments = (ins Ptr_PtrType:$base, AnySignlessIntegerOrIndex:$offset);
-  let results = (outs Ptr_PtrType:$result);
+  let arguments = (ins TT_PtrLike:$base, AnySignlessIntegerOrIndex:$offset);
+  let results = (outs TT_PtrLike:$result);
   let assemblyFormat = "$base $offset attr-dict `:` type($base) `,` type($offset) `to` type($result)";
 }
 

lib/Conversion/TritonToLinalgExperimental/CollapseShape.cpp

@@ -25,7 +25,6 @@
 #include "mlir/Dialect/Bufferization/IR/Bufferization.h"
 #include "mlir/Dialect/Linalg/IR/Linalg.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
-#include "mlir/Dialect/Ptr/IR/PtrDialect.h"
 #include "mlir/Pass/PassManager.h"
 #include "mlir/Transforms/Passes.h"
 

lib/Conversion/TritonToLinalgExperimental/ReconcilePtrCastsPass.cpp

@@ -16,20 +16,17 @@
 
 #include "mlir/Conversion/ReconcileUnrealizedCasts/ReconcileUnrealizedCasts.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
-#include "mlir/Dialect/Ptr/IR/PtrTypes.h"
+// #include "mlir/Dialect/Ptr/IR/PtrTypes.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/BuiltinDialect.h"
 #include "mlir/IR/BuiltinOps.h"
 #include "mlir/IR/BuiltinTypes.h"
 #include "mlir/IR/ValueRange.h"
 #include "mlir/Pass/PassManager.h"
 #include "mlir/Transforms/GreedyPatternRewriteDriver.h"
-#include "triton-shared/Conversion/TritonToLinalgExperimental/ReconcilePtrCasts.h"
-
-#include "triton/Dialect/Triton/IR/Types.h"
-
 #include "triton-shared/Conversion/TritonToLinalgExperimental/ReconcilePtrCasts.h"
 #include "triton-shared/Dialect/TPtr/IR/TPtrDialect.h"
+#include "triton/Dialect/Triton/IR/Types.h"
 
 using namespace mlir;
 using namespace triton;
@@ -87,17 +84,18 @@ struct FromMemrefConverter
     auto output = op.getResult(0);
     auto outType = output.getType();
 
-    if (unrankedInput && isa<triton::PointerType, ptr::PtrType>(outType)) {
+    if (unrankedInput && isa<triton::PointerType /*, ptr::PtrType*/>(outType)) {
       // from_memref only takes ranked memref, cast the unranked memref to
       // ranked memref first.
       auto rankedMemref = rewriter.create<memref::CastOp>(
           op.getLoc(), MemRefType::get({1}, unrankedInput.getElementType()),
           input);
       auto memrefToPtr = rewriter.create<tptr::FromMemrefOp>(
           op->getLoc(),
-          ptr::PtrType::get(
-              rewriter.getContext(),
-              tptr::DefaultMemorySpaceAttr::get(rewriter.getContext())),
+          // ptr::PtrType::get(
+          //     rewriter.getContext(),
+          //     tptr::DefaultMemorySpaceAttr::get(rewriter.getContext())),
+          triton::PointerType::get(unrankedInput.getElementType(), 1),
           rankedMemref);
 
       rewriter.replaceAllUsesWith(output, memrefToPtr);
@@ -123,7 +121,7 @@ struct ToMemrefConverter : public OpRewritePattern<UnrealizedConversionCastOp> {
     auto inType = input.getType();
     auto output = op.getResult(0);
     auto outUnrankedMemrefType = dyn_cast<UnrankedMemRefType>(output.getType());
-    if (isa<triton::PointerType, ptr::PtrType>(inType) &&
+    if (isa<triton::PointerType/*, ptr::PtrType*/>(inType) &&
         outUnrankedMemrefType) {
       // to_memref can only cast to ranked static shape memref, we have to cast
       // the resulting memref back to unranked

lib/Conversion/TritonToLinalgExperimental/TritonToLinalgExperimentalPass.cpp

@@ -5,7 +5,6 @@
 //
 //===----------------------------------------------------------------------===//
 
-#include "mlir/Dialect/Ptr/IR/PtrDialect.h"
 #include "triton-shared/Conversion/StructuredToMemref/StructuredToMemref.h"
 #include "triton-shared/Conversion/TritonArithToLinalg/TritonArithToLinalg.h"
 #include "triton-shared/Conversion/TritonPtrToMemref/TritonPtrToMemref.h"
@@ -46,7 +45,7 @@ class TritonToLinalgExperimentalPass
                     scf::SCFDialect, tensor::TensorDialect,
                     bufferization::BufferizationDialect, memref::MemRefDialect,
                     ttx::TritonTilingExtDialect, tts::TritonStructuredDialect,
-                    tptr::TPtrDialect, ptr::PtrDialect>();
+                    tptr::TPtrDialect>();
   }
 
   void runOnOperation() override {