Revert "removing mlir ptr dialect and ptr types in codes"

This reverts commit 920d006.

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,14 +3,11 @@
 
 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";
@@ -27,8 +24,8 @@ def TPtr_Dialect : Dialect {
     void registerTypes();
   }];
 
- let dependentDialects = [
-  "mlir::triton::TritonDialect"
+  let dependentDialects = [
+    "mlir::ptr::PtrDialect"
   ];
 
   let useDefaultAttributePrinterParser = 1;
@@ -50,10 +47,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
@@ -75,7 +72,7 @@ def TPTR_IntToPtrOp : TPTR_Op<"inttoptr", [
     ```
   }];
   let arguments = (ins AnySignlessIntegerOrIndex:$arg);
-  let results = (outs TT_PtrLike:$res);
+  let results = (outs Ptr_PtrType:$res);
   let assemblyFormat = "$arg attr-dict `:` type($arg) `to` type($res)";
 }
 
@@ -91,61 +88,45 @@ def TPTR_PtrToIntOp : TPTR_Op<"ptrtoint", [
     %int = ptr.ptrtoint %ptr : !ptr.ptr<1 : i32> to i32
     ```
   }];
-  let arguments = (ins TT_PtrLike:$arg);
+  let arguments = (ins Ptr_PtrType:$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), [{
-//       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_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_FromMemrefOp : TPTR_Op<"from_memref", [Pure]> {
   let arguments = (ins AnyMemRef:$input);
-  let results = (outs TT_PtrLike:$result);
+  let results = (outs Ptr_PtrType:$result);
   let assemblyFormat = "$input attr-dict `:` type($input) `to` type($result)";
 }
 
 def TPTR_ToMemrefOp : TPTR_Op<"to_memref", [
     Pure  ]> {
-  let arguments = (ins TT_PtrLike:$arg);
+  let arguments = (ins Ptr_PtrType:$arg);
   let results = (outs AnyStaticShapeMemRef:$res);
   let assemblyFormat = "$arg attr-dict `:` type($arg) `to` type($res)";
 }
@@ -162,8 +143,8 @@ def TPTR_PtrAddOp : TPTR_Op<"ptradd", [Pure, AllTypesMatch<["base", "result"]>]>
     ```
   }];
 
-  let arguments = (ins TT_PtrLike:$base, AnySignlessIntegerOrIndex:$offset);
-  let results = (outs TT_PtrLike:$result);
+  let arguments = (ins Ptr_PtrType:$base, AnySignlessIntegerOrIndex:$offset);
+  let results = (outs Ptr_PtrType:$result);
   let assemblyFormat = "$base $offset attr-dict `:` type($base) `,` type($offset) `to` type($result)";
 }
 

lib/Conversion/TritonToLinalgExperimental/CollapseShape.cpp

@@ -25,6 +25,7 @@
 #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,7 +16,7 @@
 
 #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"
@@ -25,9 +25,12 @@
 #include "mlir/Pass/PassManager.h"
 #include "mlir/Transforms/GreedyPatternRewriteDriver.h"
 #include "triton-shared/Conversion/TritonToLinalgExperimental/ReconcilePtrCasts.h"
-#include "triton-shared/Dialect/TPtr/IR/TPtrDialect.h"
+
 #include "triton/Dialect/Triton/IR/Types.h"
 
+#include "triton-shared/Conversion/TritonToLinalgExperimental/ReconcilePtrCasts.h"
+#include "triton-shared/Dialect/TPtr/IR/TPtrDialect.h"
+
 using namespace mlir;
 using namespace triton;
 
@@ -84,18 +87,17 @@ 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())),
-          triton::PointerType::get(unrankedInput.getElementType(), 1),
+          ptr::PtrType::get(
+              rewriter.getContext(),
+              tptr::DefaultMemorySpaceAttr::get(rewriter.getContext())),
           rankedMemref);
 
       rewriter.replaceAllUsesWith(output, memrefToPtr);
@@ -121,7 +123,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,6 +5,7 @@
 //
 //===----------------------------------------------------------------------===//
 
+#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"
@@ -45,7 +46,7 @@ class TritonToLinalgExperimentalPass
                     scf::SCFDialect, tensor::TensorDialect,
                     bufferization::BufferizationDialect, memref::MemRefDialect,
                     ttx::TritonTilingExtDialect, tts::TritonStructuredDialect,
-                    tptr::TPtrDialect>();
+                    tptr::TPtrDialect, ptr::PtrDialect>();
   }
 
   void runOnOperation() override {