[Rebase]: ttg.ViewOp->ttg.ReshapeOp in commit: 607190f… (#131)

* [Rebase]: ttg.ViewOp->ttg.ReshapeOp in commit: 607190f

* Update TODO failed cases: test_convertmma2mma

Author: quintinwang5

.github/tests/triton_todo_failure_tests.log

@@ -155,3 +155,43 @@ test_subprocess.py::test_print[device_print_large-int32]
 test_subprocess.py::test_print[device_print_multiple_args-int32]
 test_subprocess.py::test_print[print_multiple_args-int32]
 test_subprocess.py::test_print[print_no_arg-int32]
+test_core.py::test_convertmma2mma[mma_pair0-float16-64-1]
+test_core.py::test_convertmma2mma[mma_pair0-float16-1-64]
+test_core.py::test_convertmma2mma[mma_pair0-float16-64-64]
+test_core.py::test_convertmma2mma[mma_pair0-float16-128-128]
+test_core.py::test_convertmma2mma[mma_pair0-float16-256-256]
+test_core.py::test_convertmma2mma[mma_pair1-float16-64-1]
+test_core.py::test_convertmma2mma[mma_pair1-float16-1-64]
+test_core.py::test_convertmma2mma[mma_pair1-float16-64-64]
+test_core.py::test_convertmma2mma[mma_pair1-float16-128-128]
+test_core.py::test_convertmma2mma[mma_pair1-float16-256-256]
+test_core.py::test_convertmma2mma[mma_pair2-float16-64-1]
+test_core.py::test_convertmma2mma[mma_pair2-float16-1-64]
+test_core.py::test_convertmma2mma[mma_pair2-float16-64-64]
+test_core.py::test_convertmma2mma[mma_pair2-float16-128-128]
+test_core.py::test_convertmma2mma[mma_pair2-float16-256-256]
+test_core.py::test_convertmma2mma[mma_pair3-float16-64-1]
+test_core.py::test_convertmma2mma[mma_pair3-float16-1-64]
+test_core.py::test_convertmma2mma[mma_pair3-float16-64-64]
+test_core.py::test_convertmma2mma[mma_pair3-float16-128-128]
+test_core.py::test_convertmma2mma[mma_pair3-float16-256-256]
+test_core.py::test_convertmma2mma[mma_pair0-float16-64-1]
+test_core.py::test_convertmma2mma[mma_pair0-float16-1-64]
+test_core.py::test_convertmma2mma[mma_pair0-float16-64-64]
+test_core.py::test_convertmma2mma[mma_pair0-float16-128-128]
+test_core.py::test_convertmma2mma[mma_pair0-float16-256-256]
+test_core.py::test_convertmma2mma[mma_pair1-float16-64-1]
+test_core.py::test_convertmma2mma[mma_pair1-float16-1-64]
+test_core.py::test_convertmma2mma[mma_pair1-float16-64-64]
+test_core.py::test_convertmma2mma[mma_pair1-float16-128-128]
+test_core.py::test_convertmma2mma[mma_pair1-float16-256-256]
+test_core.py::test_convertmma2mma[mma_pair2-float16-64-1]
+test_core.py::test_convertmma2mma[mma_pair2-float16-1-64]
+test_core.py::test_convertmma2mma[mma_pair2-float16-64-64]
+test_core.py::test_convertmma2mma[mma_pair2-float16-128-128]
+test_core.py::test_convertmma2mma[mma_pair2-float16-256-256]
+test_core.py::test_convertmma2mma[mma_pair3-float16-64-1]
+test_core.py::test_convertmma2mma[mma_pair3-float16-1-64]
+test_core.py::test_convertmma2mma[mma_pair3-float16-64-64]
+test_core.py::test_convertmma2mma[mma_pair3-float16-128-128]
+test_core.py::test_convertmma2mma[mma_pair3-float16-256-256]

lib/Conversion/TritonGPUToSPIRV/CMakeLists.txt

@@ -1,6 +1,6 @@
 add_mlir_conversion_library(TritonGPUToSPIRV
     TritonGPUToSPIRV.cpp
-    ViewOpToSPIRV.cpp
+    ReshapeOpToSPIRV.cpp
     ElementwiseOpToSPIRV.cpp
     TritonGPUToSPIRVPass.cpp
     ConvertLayoutOpToSPIRV/SharedToDotOperandFMA.cpp

lib/Conversion/TritonGPUToSPIRV/ViewOpToSPIRV.cpp renamed to lib/Conversion/TritonGPUToSPIRV/ReshapeOpToSPIRV.cpp

@@ -1,4 +1,4 @@
-#include "ViewOpToSPIRV.h"
+#include "ReshapeOpToSPIRV.h"
 
 using namespace mlir;
 using namespace mlir::triton;
@@ -145,16 +145,35 @@ struct CatOpSPIRVConversion
   }
 };
 
-struct ViewOpSPIRVConversion : public ConvertTritonGPUOpToSPIRVPattern<ViewOp> {
-  using OpAdaptor = typename ViewOp::Adaptor;
+struct ReshapeOpSPIRVConversion
+    : public ConvertTritonGPUOpToSPIRVPattern<ReshapeOp> {
+  using OpAdaptor = typename ReshapeOp::Adaptor;
   using ConvertTritonGPUOpToSPIRVPattern<
-      ViewOp>::ConvertTritonGPUOpToSPIRVPattern;
+      ReshapeOp>::ConvertTritonGPUOpToSPIRVPattern;
 
   LogicalResult
-  matchAndRewrite(ViewOp op, OpAdaptor adaptor,
+  matchAndRewrite(ReshapeOp op, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {
     Location loc = op->getLoc();
     auto resultTy = op.getType().template cast<RankedTensorType>();
+    auto srcTy = op.getSrc().getType().template cast<RankedTensorType>();
+    if (!op.getAllowReorder()) {
+      // Only support trivial block layouts for now.
+      auto mod = op->getParentOfType<ModuleOp>();
+      int numWarps = triton::gpu::TritonGPUDialect::getNumWarps(mod);
+      int threadsPerWarp =
+          triton::gpu::TritonGPUDialect::getThreadsPerWarp(mod);
+      int numCTAs = triton::gpu::TritonGPUDialect::getNumCTAs(mod);
+      assert(resultTy.getEncoding() == triton::gpu::getDefaultBlockedEncoding(
+                                           op.getContext(), resultTy.getShape(),
+                                           numWarps, threadsPerWarp, numCTAs) &&
+             "ReshapeOp lowering only support block encoding right now.");
+      assert(srcTy.getEncoding() == triton::gpu::getDefaultBlockedEncoding(
+                                        op.getContext(), srcTy.getShape(),
+                                        numWarps, threadsPerWarp, numCTAs) &&
+             "ReshapeOp lowering only support block encoding right now.");
+    }
+
     auto vals = this->getTypeConverter()->unpackLLElements(
         loc, adaptor.getSrc(), rewriter, op.getOperand().getType());
     Value ret =
@@ -235,7 +254,7 @@ void populateViewOpToSPIRVPatterns(
     mlir::ModuleAllocation *allocation, mlir::Value smem,
     mlir::PatternBenefit benefit,
     std::map<std::string, int> &computeCapability) {
-  patterns.add<ViewOpSPIRVConversion>(typeConverter, context, benefit);
+  patterns.add<ReshapeOpSPIRVConversion>(typeConverter, context, benefit);
   patterns.add<ExpandDimsOpSPIRVConversion>(typeConverter, context, benefit);
   patterns.add<SplatOpSPIRVConversion>(typeConverter, context, benefit);
   patterns.add<ArithConstantSplatOpSPIRVConversion>(

lib/Conversion/TritonGPUToSPIRV/ViewOpToSPIRV.h renamed to lib/Conversion/TritonGPUToSPIRV/ReshapeOpToSPIRV.h

@@ -1,5 +1,5 @@
-#ifndef TRITON_VIEWOPTOSPIRV_H
-#define TRITON_VIEWOPTOSPIRV_H
+#ifndef TRITON_RESHAPEOPTOSPIRV_H
+#define TRITON_RESHAPEOPTOSPIRV_H
 
 #include "TritonGPUToSPIRVBase.h"
 #include "triton/Analysis/Membar.h"

lib/Conversion/TritonGPUToSPIRV/TritonGPUToSPIRVPass.cpp

@@ -22,10 +22,10 @@
 #include "ElementwiseOpToSPIRV.h"
 #include "LoadStoreOpToSPIRV.h"
 #include "ReduceOpToSPIRV.h"
+#include "ReshapeOpToSPIRV.h"
 #include "ScanOpToSPIRV.h"
 #include "TritonGPUToSPIRV.h"
 #include "TypeConverter.h"
-#include "ViewOpToSPIRV.h"
 
 #include "mlir/Dialect/ControlFlow/IR/ControlFlowOps.h"
 

triton_hash.txt

@@ -1 +1 @@
-f168b148ecdd067205c6066bc3e6939fd67ab893
+607190fda19daa9ecbccd7b10ed49d24b7916fd2