More progress on TritonGPUTypeConverter & TritonGPUConversionTarget

lib/Dialect/TritonGPU/CMakeLists.txt

@@ -1 +1,2 @@
 add_subdirectory(IR)
+add_subdirectory(Transforms)

lib/Dialect/TritonGPU/IR/Dialect.cpp

@@ -1,12 +1,52 @@
 #include "triton/Dialect/TritonGPU/IR/Dialect.h"
+#include "mlir/IR/DialectImplementation.h"
+#include "llvm/ADT/TypeSwitch.h"
+#include <llvm-6.0/llvm/Support/ErrorHandling.h>
 
 #include "triton/Dialect/TritonGPU/IR/Dialect.cpp.inc"
 
 using namespace mlir::triton::gpu;
 
+//===----------------------------------------------------------------------===//
+// Attribute methods
+//===----------------------------------------------------------------------===//
+#define GET_ATTRDEF_CLASSES
+#include "triton/Dialect/TritonGPU/IR/TritonGPUAttrDefs.cpp.inc"
+
+mlir::Attribute 
+TritonGPUDistributedEncodingAttr::parse(mlir::AsmParser &parser, mlir::Type type) {
+  llvm_unreachable("Not implemented");
+}
+
+void TritonGPUDistributedEncodingAttr::print(mlir::AsmPrinter &printer) const {
+  llvm_unreachable("Not implemented");
+}
+
+mlir::Attribute 
+TritonGPUMmaEncodingAttr::parse(mlir::AsmParser &parser, ::mlir::Type type) {
+  llvm_unreachable("Not implemented");
+}
+
+void TritonGPUMmaEncodingAttr::print(mlir::AsmPrinter &printer) const {
+  llvm_unreachable("Not implemented");
+}
+
+mlir::Attribute
+TritonGPUSharedEncodingAttr::parse(mlir::AsmParser &parser, ::mlir::Type type) {
+  llvm_unreachable("Not implemented");
+}
+
+void TritonGPUSharedEncodingAttr::print(mlir::AsmPrinter &printer) const {
+  llvm_unreachable("Not implemented");
+}
+
 void TritonGPUDialect::initialize() {
   addOperations<
 #define GET_OP_LIST
 #include "triton/Dialect/TritonGPU/IR/Ops.cpp.inc"
   >();
 }
+
+
+#define GET_OP_CLASSES
+#include "triton/Dialect/TritonGPU/IR/Ops.cpp.inc"

lib/Dialect/TritonGPU/Transforms/CMakeLists.txt

@@ -0,0 +1,10 @@
+add_mlir_dialect_library(TritonGPUConversion
+  TritonGPUConversion.cpp
+
+  # ADDITIONAL_HEADER_DIRS
+
+  LINK_LIBS PUBLIC
+  TritonIR
+  TritonGPUIR
+  # MLIRTransformUtils
+)

lib/Dialect/TritonGPU/Transforms/TritonGPUConversion.cpp

@@ -0,0 +1,68 @@
+#include "triton/Dialect/TritonGPU/Transforms/TritonGPUConversion.h"
+#include "triton/Dialect/Triton/IR/Dialect.h"
+#include "triton/Dialect/TritonGPU/IR/Dialect.h"
+#include <algorithm>
+
+using namespace mlir;
+
+//
+// TypeConverter
+//
+TritonGPUTypeConverter::TritonGPUTypeConverter(MLIRContext *context, 
+                                               int numThreads)
+    : context(context), numThreads(numThreads) {
+  addConversion([&](RankedTensorType tensorType) -> RankedTensorType {
+    llvm::ArrayRef<int64_t> shape = tensorType.getShape();
+    Type elementType = tensorType.getElementType();
+    int64_t rank = tensorType.getRank();
+    int64_t numElements = tensorType.getNumElements();
+
+    // TODO: we should raise exception here.
+    assert(numElements > numThreads);
+    assert(numElements % numThreads == 0);
+
+    // assert no encoding?
+
+    // Now we assume:
+    //   contiguous = 1, order = 0, 1, 2, ..., 
+    llvm::SmallVector<unsigned> threadTileSize(rank, 1); // naive layout
+    llvm::SmallVector<unsigned> blockTileSize(rank);
+    llvm::SmallVector<unsigned> order(rank);
+    int remainingThreads = numThreads;
+    for (int64_t dim = 0; dim < rank; ++dim) {
+      blockTileSize[dim] = std::clamp(remainingThreads, 1, int(shape[dim]));
+      order[dim] = dim;
+
+      remainingThreads /= blockTileSize[dim];
+      // TODO: will we need repetition?
+    }
+    Attribute encoding = triton::gpu::TritonGPUDistributedEncodingAttr::get(
+        context, threadTileSize, blockTileSize, order);
+    return RankedTensorType::get(shape, elementType, encoding);
+  });
+}
+
+//
+// TritonGPUConversion
+//
+TritonGPUConversionTarget::TritonGPUConversionTarget(MLIRContext &context)
+    : ConversionTarget(context) {
+  addLegalDialect<triton::TritonDialect,
+                  arith::ArithmeticDialect,
+                  scf::SCFDialect>();
+
+  // Some ops from SCF are illegal
+  addIllegalOp<scf::ExecuteRegionOp, scf::ParallelOp, 
+               scf::ReduceOp, scf::ReduceReturnOp>();
+
+  // // We have requirements for the data layouts
+  // addDynamicallyLegalOp<triton::DotOp>([](triton::DotOp dotOp) -> bool {
+  //   Attribute aEncoding = dotOp.a().getType().cast<RankedTensorType>().getEncoding();
+  //   Attribute bEncoding = dotOp.b().getType().cast<RankedTensorType>().getEncoding();
+  //   if (aEncoding && aEncoding.isa<triton::gpu::TritonGPUSharedEncodingAttr>() &&
+  //       bEncoding && bEncoding.isa<triton::gpu::TritonGPUSharedEncodingAttr>())
+  //     return true;
+  //   return false;
+  // });
+
+}