add tptr-to-llvm pass

Author: Realtyxxx

CMakeLists.txt

@@ -40,6 +40,7 @@ if(TRITON_SHARED_BUILD_CPU_BACKEND)
         MLIRControlFlowToLLVM
         MLIRMemRefToLLVM
         MLIRVectorToSCF
+        TPtrToLLVM
         MLIRUBToLLVM
         # MLIRMathToLibm
 

backend/compiler.py

@@ -100,6 +100,8 @@ def _ttsharedir_to_llir(ttsharedir: str):
         triton_shared.to_llir.add_memref_expand(pm)
         triton_shared.to_llir.add_finalize_memref_to_llvm(pm)
         triton_shared.to_llir.add_convert_func_to_llvm(pm)
+        # triton_shared.debug.enable_mlir_debug("tptr-to-llvm")
+        triton_shared.to_llir.add_convert_tptr_to_llvm(pm)
 
         triton_shared.to_llir.add_convert_cf_to_llvm(pm)
         triton_shared.to_llir.add_lower_affine(pm)

include/triton-shared/Conversion/CMakeLists.txt

@@ -6,3 +6,4 @@ add_subdirectory(TritonPtrToMemref)
 add_subdirectory(TritonToUnstructured)
 add_subdirectory(StructuredToMemref)
 add_subdirectory(UnstructuredToMemref)
+add_subdirectory(TPtrToLLVM)

include/triton-shared/Conversion/TPtrToLLVM/CMakeLists.txt

@@ -0,0 +1,3 @@
+set(LLVM_TARGET_DEFINITIONS Passes.td)
+mlir_tablegen(Passes.h.inc -gen-pass-decls --name TPtrToLLVM)
+add_public_tablegen_target(TPtrToLLVMConversionPassIncGen)

include/triton-shared/Conversion/TPtrToLLVM/Passes.h

@@ -0,0 +1,15 @@
+#ifndef TPTR_TO_LLVM_CONVERSION_PASSES_H
+#define TPTR_TO_LLVM_CONVERSION_PASSES_H
+
+#include "triton-shared/Conversion/TPtrToLLVM/TPtrToLLVM.h"
+
+namespace mlir {
+namespace tptr {
+
+#define GEN_PASS_REGISTRATION
+#include "triton-shared/Conversion/TPtrToLLVM/Passes.h.inc"
+
+}  // namespace triton
+}  // namespace mlir
+
+#endif

include/triton-shared/Conversion/TPtrToLLVM/Passes.td

@@ -0,0 +1,20 @@
+#ifndef TPTR_TO_LLVM_CONVERSION_PASSES
+#define TPTR_TO_LLVM_CONVERSION_PASSES
+
+include "mlir/Pass/PassBase.td"
+
+def TPtrToLLVM : Pass<"tptr-to-llvm", "mlir::ModuleOp"> {
+  let summary = "Convert Tptr operations into LLVM";
+  let options = [
+//       Option<"addptrToLLVM", "ptradd-to-llvm", "bool", /*default*/"true",
+//              "Convert tptr.ptradd on tensors to llvm">,
+//       Option<"fromMemrefToLLVM", "from-memref-to-llvm", "bool", /*default*/"true",
+//              "Convert tptr.from_memref on tensors to llvm">,
+//       Option<"toMemrefToLLVM", "to-memref-to-llvm", "bool", /*default*/"true",
+//              "Convert tptr.to_memref on tensors to llvm">,
+//       Option<"typeoffsetToConst", "typeoffset-to-const", "bool", /*default*/"true",
+//              "Convert tptr.typeoffset to llvm.mlir.const">,
+  ];
+  let dependentDialects = ["mlir::tptr::TPtrDialect", "mlir::LLVM::LLVMDialect"];
+}
+#endif

include/triton-shared/Conversion/TPtrToLLVM/TPtrToLLVM.h

@@ -0,0 +1,32 @@
+#ifndef TRITON_CONVERSION_TPTR_TO_LLVM_TPTRTOLLVM_H
+#define TRITON_CONVERSION_TPTR_TO_LLVM_TPTRTOLLVM_H
+
+#include "mlir/Dialect/LLVMIR/LLVMDialect.h"
+#include "mlir/Pass/Pass.h"
+#include "mlir/Transforms/DialectConversion.h"
+#include "triton-shared/Dialect/TPtr/IR/TPtrDialect.h"
+
+namespace mlir {
+namespace tptr {
+
+#define GEN_PASS_DECL
+// #define GEN_PASS_CLASSES
+// #define GEN_PASS_DECL_TPTRTOLLVM
+// #define GEN_PASS_DEF_TPTRTOLLVM
+#include "triton-shared/Conversion/TPtrToLLVM/Passes.h.inc"
+
+// void populateTPtrToLLVMCanonicalizationPatterns(RewritePatternSet &patterns);
+
+void populateTPtrToLLVMConversionPatterns(RewritePatternSet &patterns,
+                                          TypeConverter &typeconverter);
+
+std::unique_ptr<OperationPass<ModuleOp>> createTPtrToLLVMPass();
+
+static bool isOneToOneCast(UnrealizedConversionCastOp op) {
+  return (op.getInputs().size() == 1 && op->getNumResults() == 1);
+}
+
+}  // namespace tptr
+}  // namespace mlir
+
+#endif

lib/Conversion/CMakeLists.txt

@@ -6,3 +6,4 @@ add_subdirectory(TritonArithToLinalg)
 add_subdirectory(StructuredToMemref)
 add_subdirectory(TritonPtrToMemref)
 add_subdirectory(UnstructuredToMemref)
+add_subdirectory(TPtrToLLVM)

lib/Conversion/TPtrToLLVM/CMakeLists.txt

@@ -0,0 +1,17 @@
+add_triton_library(TPtrToLLVM
+  TPtrToLLVMPass.cpp
+  TPtrToLLVM.cpp
+
+  DEPENDS
+  TPtrToLLVMConversionPassIncGen
+  TPtrTableGen
+
+  LINK_LIBS PUBLIC
+  MLIRIR
+  MLIRPass
+  MLIRTransforms
+  MLIRSupport
+  MLIRReconcileUnrealizedCasts
+  TPtrIR
+  MLIRDialectUtils
+)