[Gluon] Add AutoLayout for backward layout inference (#7447)

This adds `ttgl.AutoLayout` which lowers to the new
`gluon::AutoEncodingAttr` encoding.

In the language, we allow this encoding to propagate through forward
type inference using a custom layout inference interface which always
returns auto regardless of the operation. So e.g. `sum(a, 1)` where
`a.layout` is auto layout will return auto layout, not `SliceLayout(1,
AutoLayout())`.

This also adds the gluon-specific `--gluon-resolve-auto-encodings` pass,
which finds layout conversion from auto encoding to a concrete encoding
and propagates the concrete encoding through the graph.

The pass will error out if there are any conflicting inferences, rather
than duplicating computations as the backward rematerialization pass
does. It also errors out if the inference fails for any reason. Note
that the errors are really bad at the moment. Will need to fix in a
follow-up.

Author: peterbell10

bin/RegisterTritonDialects.h

@@ -4,6 +4,7 @@
 #include "third_party/nvidia/include/Dialect/NVGPU/IR/Dialect.h"
 #include "third_party/nvidia/include/Dialect/NVWS/IR/Dialect.h"
 #include "third_party/proton/dialect/include/Dialect/Proton/IR/Dialect.h"
+#include "triton/Dialect/Gluon/Transforms/Passes.h"
 #include "triton/Dialect/Triton/IR/Dialect.h"
 #include "triton/Dialect/TritonGPU/IR/Dialect.h"
 #include "triton/Dialect/TritonInstrument/IR/Dialect.h"
@@ -50,6 +51,7 @@ inline void registerTritonDialects(mlir::DialectRegistry &registry) {
   mlir::triton::gpu::registerTritonGPUPasses();
   mlir::triton::nvidia_gpu::registerTritonNvidiaGPUPasses();
   mlir::triton::instrument::registerTritonInstrumentPasses();
+  mlir::triton::gluon::registerGluonPasses();
   mlir::test::registerTestAliasPass();
   mlir::test::registerTestAlignmentPass();
   mlir::test::registerAMDTestAlignmentPass();
@@ -105,5 +107,6 @@ inline void registerTritonDialects(mlir::DialectRegistry &registry) {
       mlir::gpu::GPUDialect, mlir::LLVM::LLVMDialect, mlir::NVVM::NVVMDialect,
       mlir::triton::nvgpu::NVGPUDialect, mlir::triton::nvws::NVWSDialect,
       mlir::triton::amdgpu::TritonAMDGPUDialect,
-      mlir::triton::proton::ProtonDialect, mlir::ROCDL::ROCDLDialect>();
+      mlir::triton::proton::ProtonDialect, mlir::ROCDL::ROCDLDialect,
+      mlir::triton::gluon::GluonDialect>();
 }

include/triton/Dialect/CMakeLists.txt

@@ -2,3 +2,4 @@ add_subdirectory(Triton)
 add_subdirectory(TritonGPU)
 add_subdirectory(TritonNvidiaGPU)
 add_subdirectory(TritonInstrument)
+add_subdirectory(Gluon)

include/triton/Dialect/Gluon/CMakeCache.txt

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

include/triton/Dialect/Gluon/CMakeLists.txt

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

include/triton/Dialect/Gluon/IR/CMakeLists.txt

@@ -0,0 +1,16 @@
+set(MLIR_BINARY_DIR ${CMAKE_BINARY_DIR})
+
+set(LLVM_TARGET_DEFINITIONS GluonDialect.td)
+mlir_tablegen(Dialect.h.inc -gen-dialect-decls -dialect=gluon)
+mlir_tablegen(Dialect.cpp.inc -gen-dialect-defs -dialect=gluon)
+mlir_tablegen(Ops.h.inc -gen-op-decls)
+mlir_tablegen(Ops.cpp.inc -gen-op-defs)
+add_mlir_doc(GluonDialect GluonDialect dialects/ -gen-dialect-doc)
+add_public_tablegen_target(GluonTableGen)
+
+set(LLVM_TARGET_DEFINITIONS GluonAttrDefs.td)
+mlir_tablegen(GluonAttrDefs.h.inc -gen-attrdef-decls)
+mlir_tablegen(GluonAttrDefs.cpp.inc -gen-attrdef-defs)
+mlir_tablegen(OpsEnums.h.inc -gen-enum-decls)
+mlir_tablegen(OpsEnums.cpp.inc -gen-enum-defs)
+add_public_tablegen_target(GluonAttrDefsIncGen)

include/triton/Dialect/Gluon/IR/Dialect.h

@@ -0,0 +1,8 @@
+#pragma once
+#include "triton/Dialect/Triton/IR/Dialect.h"
+#include "triton/Dialect/TritonGPU/IR/Dialect.h"
+
+#include "triton/Dialect/Gluon/IR/Dialect.h.inc"
+
+#define GET_ATTRDEF_CLASSES
+#include "triton/Dialect/Gluon/IR/GluonAttrDefs.h.inc"

include/triton/Dialect/Gluon/IR/GluonAttrDefs.td

@@ -0,0 +1,15 @@
+#ifndef GLUON_ATTRDEFS
+#define GLUON_ATTRDEFS
+
+include "mlir/IR/AttrTypeBase.td"
+include "triton/Dialect/Gluon/IR/GluonDialect.td"
+
+def Gluon_AutoEncodingAttr : AttrDef<Gluon_Dialect, "AutoEncoding"> {
+  let mnemonic = "auto_encoding";
+  let attrName = "gluon.auto_encoding";
+  let description = [{
+    An encoding that is inferred from neighboring ops in the graph.
+  }];
+}
+
+#endif

include/triton/Dialect/Gluon/IR/GluonDialect.td

@@ -0,0 +1,22 @@
+#ifndef GLUON_DIALECT
+#define GLUON_DIALECT
+
+include "mlir/IR/OpBase.td"
+
+def Gluon_Dialect : Dialect {
+  let name = "gluon";
+  let cppNamespace = "::mlir::triton::gluon";
+  let description = [{
+    Gluon dialect.
+  }];
+
+  let dependentDialects = [
+    "triton::TritonDialect",
+    "triton::gpu::TritonGPUDialect",
+    "mlir::gpu::GPUDialect",
+  ];
+  let useDefaultAttributePrinterParser = 1;
+  let usePropertiesForAttributes = 1;
+}
+
+#endif

include/triton/Dialect/Gluon/Transforms/CMakeLists.txt

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

include/triton/Dialect/Gluon/Transforms/Passes.h

@@ -0,0 +1,13 @@
+#pragma once
+#include "mlir/IR/BuiltinOps.h"
+#include "mlir/Pass/Pass.h"
+#include "triton/Dialect/Gluon/IR/Dialect.h"
+#include <memory>
+
+namespace mlir::triton::gluon {
+
+#define GEN_PASS_DECL
+#define GEN_PASS_REGISTRATION
+#include "triton/Dialect/Gluon/Transforms/Passes.h.inc"
+
+} // namespace mlir::triton::gluon