Merge triton-mlir branch - Complete rewrite of the backend from scratch (#1004)

This PR merges the `triton-mlir` branch, in which we have been quietly
rewriting the Triton backend from scratch to increase maintainability,
stability and ultimately performance. Changes to the runtime are
minimal, and this new version aims to remain backward-compatible with
the previous commit. The legacy backend is now officially deprecated,
but can still be accessed via the `legacy-backend` tag.

Co-authored-by: Keren Zhou <kerenzhou@openai.com>
Co-authored-by: Yan Chunwei <yanchunwei@outlook.com>
Co-authored-by: goostavz <109190422+goostavz@users.noreply.github.com>
Co-authored-by: Shintaro Iwasaki <siwasaki@fb.com>
Co-authored-by: Yan Da <dyanab@connect.ust.hk>
Co-authored-by: Jun Yang <yangjunpro@gmail.com>
Co-authored-by: Ian Bearman <ianb@microsoft.com>
Co-authored-by: Jason Ansel <jansel@jansel.net>
Co-authored-by: Qingyi Liu <qingyil@nvidia.com>
Co-authored-by: ben-zhang-609 <110140741+ben-zhang-609@users.noreply.github.com>
Co-authored-by: Chenggang Zhao <lyricz@yeah.net>
Co-authored-by: ben-zhang-609 <benzh609@gmail.com>
Co-authored-by: dongdongl <dongdongl@nvidia.com>

unittest/Analysis/CMakeLists.txt

@@ -0,0 +1,5 @@
+add_triton_ut(
+  NAME TestTritonAnalysis
+  SRCS UtilityTest.cpp
+  LIBS TritonAnalysis
+)

unittest/Analysis/UtilityTest.cpp

@@ -0,0 +1,29 @@
+//===- UtilityTest.cpp - Tests for
+// Utility----------------------------------===//
+//
+//===----------------------------------------------------------------------===//
+
+#include "triton/Analysis/Utility.h"
+#include <gtest/gtest.h>
+
+namespace mlir {
+
+TEST(Analysis, reorder) {
+  SmallVector<int> shape({10, 20, 30});
+  {
+    SmallVector<unsigned> order({2, 1, 0});
+    auto reordered = reorder<int>(shape, order);
+    EXPECT_EQ(reordered[0], 30);
+    EXPECT_EQ(reordered[1], 20);
+    EXPECT_EQ(reordered[2], 10);
+  }
+  {
+    SmallVector<unsigned> order({1, 0, 2});
+    auto reordered = reorder<int>(shape, order);
+    EXPECT_EQ(reordered[0], 20);
+    EXPECT_EQ(reordered[1], 10);
+    EXPECT_EQ(reordered[2], 30);
+  }
+}
+
+} // namespace mlir

unittest/CMakeLists.txt

@@ -0,0 +1,29 @@
+
+include (${CMAKE_CURRENT_SOURCE_DIR}/googletest.cmake)
+
+include(GoogleTest)
+enable_testing()
+
+function(add_triton_ut)
+  set(options)
+  set(oneValueArgs NAME)
+  set(multiValueArgs SRCS LIBS)
+  cmake_parse_arguments(_ "${options}" "${oneValueArgs}" "${multiValueArgs}" ${ARGN})
+  add_test(NAME ${__NAME}
+          COMMAND ${__NAME})
+  add_executable(
+          ${__NAME}
+          ${__SRCS})
+  target_link_libraries(
+          ${__NAME}
+          PRIVATE
+          GTest::gtest_main
+          gmock
+          ${__LIBS})
+
+  gtest_discover_tests(${__NAME})
+endfunction()
+
+add_subdirectory(Analysis)
+add_subdirectory(Conversion)
+add_subdirectory(Dialect)

unittest/Conversion/CMakeLists.txt

@@ -0,0 +1 @@
+add_subdirectory(TritonGPUToLLVM)

unittest/Conversion/TritonGPUToLLVM/CMakeLists.txt

@@ -0,0 +1,5 @@
+add_triton_ut(
+	NAME TestPtxAsmFormat
+	SRCS PTXAsmFormatTest.cpp
+	LIBS TritonGPUToLLVM
+)

unittest/Conversion/TritonGPUToLLVM/PTXAsmFormatTest.cpp

@@ -0,0 +1,147 @@
+#include "triton/Conversion/TritonGPUToLLVM/PTXAsmFormat.h"
+#include "mlir/Dialect/Arithmetic/IR/Arithmetic.h"
+#include "mlir/IR/Builders.h"
+#include "triton/Dialect/Triton/IR/Dialect.h"
+
+#include <gtest/gtest.h>
+
+namespace mlir {
+namespace triton {
+class PTXAsmFormatTest : public ::testing::Test {
+protected:
+  static constexpr int numValues = 4;
+
+  PTXAsmFormatTest() {
+    ctx.loadDialect<arith::ArithmeticDialect>();
+
+    createValues();
+  }
+
+  // Creates the test values.
+  void createValues() {
+    OpBuilder builder(&ctx);
+    builder.setInsertionPointToStart(&block);
+
+    // a b1 value for predicate.
+    v[0] = builder.create<arith::ConstantIntOp>(builder.getUnknownLoc(), 1, 1);
+    for (int i = 0; i < numValues; i++) {
+      v[i + 1] =
+          builder.create<arith::ConstantIntOp>(builder.getUnknownLoc(), i, 32);
+    }
+  }
+
+  MLIRContext ctx;
+  Block block;
+  Value v[numValues + 1];
+};
+
+TEST_F(PTXAsmFormatTest, basic) {
+  PTXBuilder builder;
+
+  // Create the operands needed by the instructions in the PTX code.
+  auto *cst = builder.newConstantOperand(1);
+  auto *val = builder.newOperand(v[1], "=r");
+
+  // create an instruction
+  auto &mov = *builder.create("mov.b16");
+
+  mov(val, cst).predicate(v[0]);
+  ASSERT_EQ(builder.dump(), "@$1 mov.b16 $0, 0x1;");
+
+  auto values = builder.getAllMLIRArgs();
+  ASSERT_EQ(values[0], v[1]); // $0 -> v[1]
+  ASSERT_EQ(values[1], v[0]); // $1 -> v[0]
+
+  auto constraints = builder.getConstraints();
+  ASSERT_EQ(constraints, "=r,b"); // $0 -> =r, $1 -> b
+}
+
+TEST_F(PTXAsmFormatTest, complexInstruction) {
+  using triton::CacheModifier;
+  using triton::EvictionPolicy;
+
+  PTXBuilder builder;
+
+  int width = 16;
+  int nWords = 2;
+
+  Value predicateVal = v[0];
+  Value addrVal = v[1];
+
+  auto addr = builder.newAddrOperand(addrVal, "l", 128 /*offset*/);
+
+  bool isVolatile = false;
+  auto cache = triton::CacheModifier::CA;
+  auto cachePriority = triton::EvictionPolicy::EVICT_FIRST;
+  bool hasL2EvictPolicy = true;
+
+  auto &ld =
+      builder
+          .create<>("ld") //
+          ->o("volatile", isVolatile)
+          .global()
+          .o("ca", cache == CacheModifier::CA)
+          .o("cg", cache == CacheModifier::CG)
+          .o("L1::evict_first", cachePriority == EvictionPolicy::EVICT_FIRST)
+          .o("L1::evict_last", cachePriority == EvictionPolicy::EVICT_LAST)
+          .o("L1::cache_hint", hasL2EvictPolicy)
+          .v(nWords)
+          .b(width);
+
+  // Link the instruction to operands
+  ld(addr).predicate(predicateVal);
+
+  EXPECT_EQ(
+      builder.dump(),
+      "@$1 ld.global.ca.L1::evict_first.L1::cache_hint.v2.b16 [ $0 + 128 ];");
+  auto values = builder.getAllMLIRArgs();
+  EXPECT_EQ(values[0], addrVal);      // $0 -> predicate
+  EXPECT_EQ(values[1], predicateVal); // $1 -> addr
+  EXPECT_EQ(builder.getConstraints(), "l,b");
+}
+
+TEST_F(PTXAsmFormatTest, MultiLinePTX) {
+  PTXBuilder builder;
+
+  auto *constVal = builder.newConstantOperand(1);
+  auto *valVal0 = builder.newOperand(v[1], "=r");
+  auto *valVal1 = builder.newOperand(v[2], "=r");
+
+  auto &mov = *builder.create("mov");
+
+  mov(valVal0, constVal);
+  mov(valVal1, constVal);
+  mov(valVal1, valVal0);
+
+  EXPECT_EQ(builder.dump(), "mov $0, 0x1;\n\t"
+                            "mov $1, 0x1;\n\t"
+                            "mov $1, $0;");
+
+  auto values = builder.getAllMLIRArgs();
+  EXPECT_EQ(values[0], v[1]); // $0 -> v[1]
+  EXPECT_EQ(values[1], v[2]); // $1 -> v[2]
+}
+
+TEST_F(PTXAsmFormatTest, onlyAttachMLIRArgs) {
+  PTXBuilder builder;
+  const char *ptxCode =
+      ".param .b64 param0;\n" // prepare param0 (format string)
+      "st.param.b64 [param0], %0;\n"
+      "st.param.b64 [param0], %1;\n"
+      "st.param.b64 [param0], %2;\n";
+
+  auto &ptxSnippet = *builder.create(ptxCode);
+  auto *opr0 = builder.newOperand(v[0], "r");
+  auto *opr1 = builder.newOperand(v[1], "r");
+  auto *opr2 = builder.newOperand(v[2], "r");
+  ptxSnippet({opr1, opr2, opr0}, true);
+
+  EXPECT_EQ(builder.dump(), ptxCode);
+  ASSERT_EQ(builder.getAllMLIRArgs()[0], v[1]);
+  ASSERT_EQ(builder.getAllMLIRArgs()[1], v[2]);
+  ASSERT_EQ(builder.getAllMLIRArgs()[2], v[0]);
+  ASSERT_EQ(builder.getAllMLIRArgs().size(), 3);
+}
+
+} // namespace triton
+} // namespace mlir

unittest/Dialect/CMakeLists.txt

@@ -0,0 +1 @@
+add_subdirectory(TritonGPU)

unittest/Dialect/TritonGPU/CMakeLists.txt

@@ -0,0 +1,6 @@
+
+add_triton_ut(
+	NAME TestSwizzling
+	SRCS SwizzleTest.cpp
+	LIBS TritonGPUIR  ${dialect_libs} ${conversion_libs}
+)

unittest/Dialect/TritonGPU/SwizzleTest.cpp

@@ -0,0 +1,53 @@
+#include "triton/Dialect/TritonGPU/IR/Dialect.h"
+#include <gtest/gtest.h>
+
+using namespace mlir;
+using mlir::triton::gpu::SharedEncodingAttr;
+
+struct swizzleParams {
+  int vec;
+  int perPhase;
+  int maxPhase;
+};
+
+struct ParamT {
+  std::array<int64_t, 2> shape;
+  int opIdx;
+  int typeWidth;
+  swizzleParams refSwizzle;
+};
+
+class SwizzleDotOperandTestFixture : public ::testing::TestWithParam<ParamT> {
+protected:
+  ParamType param;
+};
+
+TEST_P(SwizzleDotOperandTestFixture, DotOperands) {
+  auto params = GetParam();
+  // init context
+  MLIRContext ctx;
+  ctx.loadDialect<triton::gpu::TritonGPUDialect>();
+  // create encoding
+  auto parent = triton::gpu::MmaEncodingAttr::get(&ctx, 2, 0, {1, 1});
+  auto encoding =
+      triton::gpu::DotOperandEncodingAttr::get(&ctx, params.opIdx, parent);
+
+  // create element type
+  Type eltType = IntegerType::get(&ctx, params.typeWidth);
+  auto layout =
+      SharedEncodingAttr::get(&ctx, encoding, params.shape, {1, 0}, eltType);
+
+  ASSERT_EQ(layout.getVec(), params.refSwizzle.vec);
+  ASSERT_EQ(layout.getPerPhase(), params.refSwizzle.perPhase);
+  ASSERT_EQ(layout.getMaxPhase(), params.refSwizzle.maxPhase);
+}
+
+INSTANTIATE_TEST_SUITE_P(TestDotOperands, SwizzleDotOperandTestFixture,
+                         ::testing::Values(ParamT{{128, 64}, 0, 16, {8, 1, 8}},
+                                           ParamT{{64, 256}, 1, 16, {8, 1, 8}},
+                                           ParamT{{128, 32}, 0, 16, {8, 2, 4}},
+                                           ParamT{{32, 128}, 1, 16, {8, 1, 8}},
+                                           ParamT{{32, 32}, 0, 16, {8, 2, 4}},
+                                           ParamT{{32, 32}, 1, 16, {8, 2, 4}},
+                                           ParamT{{16, 16}, 0, 16, {8, 4, 2}},
+                                           ParamT{{16, 16}, 1, 16, {8, 4, 2}}));

unittest/googletest.cmake

@@ -0,0 +1,23 @@
+include(FetchContent)
+
+set(GOOGLETEST_DIR "" CACHE STRING "Location of local GoogleTest repo to build against")
+
+if(GOOGLETEST_DIR)
+  set(FETCHCONTENT_SOURCE_DIR_GOOGLETEST ${GOOGLETEST_DIR} CACHE STRING "GoogleTest source directory override")
+endif()
+
+FetchContent_Declare(
+  googletest
+  GIT_REPOSITORY https://github.com/google/googletest.git
+  GIT_TAG release-1.12.1
+  )
+
+FetchContent_GetProperties(googletest)
+
+if(NOT googletest_POPULATED)
+  FetchContent_Populate(googletest)
+  if (MSVC)
+    set(gtest_force_shared_crt ON CACHE BOOL "" FORCE)
+  endif()
+  add_subdirectory(${googletest_SOURCE_DIR} ${googletest_BINARY_DIR} EXCLUDE_FROM_ALL)
+endif()