Skip to content

Commit 1cc902f

Browse files
huatuolihuatuoli
committed
[TLE][ILUVATAR] Add TLE support for alloc, local_ptr, copy, extract_tile and insert_tile.
1 parent 31d3e10 commit 1cc902f

41 files changed

Lines changed: 3717 additions & 5 deletions

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

‎CMakeLists.txt‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -38,6 +38,8 @@ if(NOT FLAGTREE_BACKEND)
3838
elseif(FLAGTREE_BACKEND STREQUAL "iluvatar")
3939
add_definitions(-D__ILUVATAR__)
4040
set(FLAGTREE_TLE OFF)
41+
set(FLAGTREE_ILUVATAR_TLE ON)
42+
add_definitions(-D__ILUVATAR_TLE__)
4143
remove_definitions(-D__TLE__)
4244
list(REMOVE_ITEM LLVM_TABLEGEN_FLAGS -D__TLE__)
4345
elseif(FLAGTREE_BACKEND STREQUAL "mthreads")
@@ -360,6 +362,10 @@ elseif(FLAGTREE_BACKEND STREQUAL "iluvatar")
360362
include_directories(${TRITON_CORE_BINARY_DIR}/include)
361363
include_directories(${TRITON_CORE_SOURCE_DIR}/backend/include)
362364
include_directories(${TRITON_CORE_BINARY_DIR}/backend/include)
365+
if(FLAGTREE_ILUVATAR_TLE)
366+
include_directories(${TRITON_CORE_SOURCE_DIR}/tle/include)
367+
include_directories(${TRITON_CORE_BINARY_DIR}/tle/include)
368+
endif()
363369
add_subdirectory(${TRITON_CORE_SOURCE_DIR}/include ${TRITON_CORE_BINARY_DIR}/include)
364370
add_subdirectory(${TRITON_CORE_SOURCE_DIR}/lib ${TRITON_CORE_BINARY_DIR}/lib)
365371
elseif(NOT FLAGTREE_BACKEND)

‎python/triton/experimental/tle/language/gpu/core.py‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@
55
from enum import Enum
66
from . import types as tle
77
from .mthreads import copy as mthreads_copy
8+
from .iluvatar import copy as iluvatar_copy
89
from triton.compiler.code_generator import flatten_values_to_ir, unflatten_ir_values
910

1011
from triton.language.core import (
@@ -362,6 +363,7 @@ def copy(
362363
tle.copy(tma_desc, local_buf, [64, 64], [x_offset, y_offset])
363364
"""
364365
mthreads_enabled = mthreads_copy.enabled()
366+
iluvatar_enabled = iluvatar_copy.enabled()
365367

366368
def normcopy(
367369
src: tl.tensor,
@@ -394,7 +396,7 @@ def normcopy(
394396
try:
395397
if direction == CopyDirection.GM_TO_LOCAL:
396398
# None fills the FlagTree hints slot; TLE copy has no hints to pass.
397-
load_extra_args = () if mthreads_enabled else (None, )
399+
load_extra_args = () if (mthreads_enabled or iluvatar_enabled) else (None, )
398400
tt_load = _semantic.load(src, mask, other, boundary_check, padding_option, cache_modifier,
399401
eviction_policy, volatile, *load_extra_args)
400402
local_ptrs = local_ptr(dst, _make_full_indices(dst, _semantic), _semantic=_semantic)

‎python/triton/experimental/tle/language/gpu/iluvatar/__init__.py‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,3 @@
1+
from . import copy
2+
3+
__all__ = ["copy"]

‎python/triton/experimental/tle/language/gpu/iluvatar/copy.py‎

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,18 @@
1+
import os
2+
3+
try:
4+
from triton._flagtree_backend import FLAGTREE_BACKEND
5+
except ModuleNotFoundError:
6+
FLAGTREE_BACKEND = os.environ.get("FLAGTREE_BACKEND", "")
7+
8+
9+
def _has_iluvatar_libtriton() -> bool:
10+
try:
11+
from triton._C import libtriton
12+
except ImportError:
13+
return False
14+
return hasattr(libtriton, "iluvatar")
15+
16+
17+
def enabled() -> bool:
18+
return FLAGTREE_BACKEND == "iluvatar" or _has_iluvatar_libtriton()

‎third_party/iluvatar/CMakeLists.txt‎

Lines changed: 22 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,12 +1,33 @@
11
include_directories(${CMAKE_CURRENT_SOURCE_DIR}/backend/include)
22
include_directories(${CMAKE_CURRENT_BINARY_DIR}/backend/include)
3+
if(FLAGTREE_ILUVATAR_TLE)
4+
include_directories(${CMAKE_CURRENT_SOURCE_DIR}/tle/include)
5+
include_directories(${CMAKE_CURRENT_BINARY_DIR}/tle/include)
6+
add_subdirectory(tle)
7+
endif()
38
add_subdirectory(backend/include)
49
add_subdirectory(backend/lib)
510
if(TRITON_BUILD_PYTHON_MODULE)
611
find_package(LLD REQUIRED CONFIG PATHS "${LLD_DIR}" NO_DEFAULT_PATH)
712
include_directories(${LLD_INCLUDE_DIRS})
813
message(STATUS "Found LLD distro-package @ ${LLD_DIR} and LLD include dirs @ ${LLD_INCLUDE_DIRS}")
9-
add_triton_plugin(TritonILUVATAR ${CMAKE_CURRENT_SOURCE_DIR}/triton_iluvatar.cc LINK_LIBS TritonILUVATARGPUToLLVM)
14+
if(FLAGTREE_ILUVATAR_TLE)
15+
set(_ILUVATAR_TLE_PLUGIN_SOURCES
16+
${CMAKE_CURRENT_SOURCE_DIR}/tle/triton_iluvatar_tle.cc)
17+
set(_ILUVATAR_TLE_PLUGIN_LIBS IluvatarTleIR IluvatarTleTransforms)
18+
set(_ILUVATAR_TLE_PLUGIN_DEPS IluvatarTleTableGen IluvatarTleTransformsIncGen)
19+
else()
20+
set(_ILUVATAR_TLE_PLUGIN_SOURCES "")
21+
set(_ILUVATAR_TLE_PLUGIN_LIBS "")
22+
set(_ILUVATAR_TLE_PLUGIN_DEPS "")
23+
endif()
24+
add_triton_plugin(TritonILUVATAR
25+
${CMAKE_CURRENT_SOURCE_DIR}/triton_iluvatar.cc
26+
${_ILUVATAR_TLE_PLUGIN_SOURCES}
27+
LINK_LIBS TritonILUVATARGPUToLLVM ${_ILUVATAR_TLE_PLUGIN_LIBS})
28+
if(FLAGTREE_ILUVATAR_TLE)
29+
add_dependencies(TritonILUVATAR ${_ILUVATAR_TLE_PLUGIN_DEPS})
30+
endif()
1031
target_link_libraries(TritonILUVATAR PRIVATE Python3::Module pybind11::headers lldCommon lldELF)
1132
endif()
1233
if(TRITON_BUILD_UT)

‎third_party/iluvatar/backend/compiler.py‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -263,6 +263,10 @@ def make_ttgir(mod, metadata, opt, capability):
263263
passes.ttgpuir.add_f32_dot_tc(pm, emuTF32)
264264
passes.ttgpuir.add_remove_layout_conversions(pm)
265265
passes.ttgpuir.add_optimize_thread_locality(pm)
266+
if hasattr(iluvatar.passes, "tle"):
267+
iluvatar.passes.tle.add_insert_local_pointer_barriers(pm)
268+
iluvatar.passes.tle.add_optimize_local_pointer_loads(pm)
269+
iluvatar.passes.tle.add_optimize_local_pointer_stores(pm)
266270
iluvatar.passes.ttgpuir.add_accelerate_matmul(pm, opt.use_sme)
267271
passes.ttgpuir.add_remove_layout_conversions(pm)
268272
iluvatar.passes.ttgpuir.add_mma_reduce_thread_locality(pm)

‎third_party/iluvatar/backend/lib/TritonILUVATARGPUToLLVM/CMakeLists.txt‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,9 @@
1+
if(FLAGTREE_ILUVATAR_TLE)
2+
set(_ILUVATAR_TLE_LIBS IluvatarTleToLLVM)
3+
else()
4+
set(_ILUVATAR_TLE_LIBS "")
5+
endif()
6+
17
add_triton_library(TritonILUVATARGPUToLLVM
28
AtomicRMWOpsEmitter.cpp
39
MaskedOpsToLLVM.cpp
@@ -19,6 +25,7 @@ add_triton_library(TritonILUVATARGPUToLLVM
1925
TritonGPUToLLVM
2026
TritonILUVATARGPUIR
2127
TritonILUVATARUtils
28+
${_ILUVATAR_TLE_LIBS}
2229
MLIRUBToLLVM
2330
LLVMCore
2431
LLVMPasses

‎third_party/iluvatar/backend/lib/TritonILUVATARGPUToLLVM/TritonGPUToLLVM.cpp‎

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,9 @@
11
#include "TritonILUVATARGPUToLLVM/Passes.h"
22

3+
#ifdef __ILUVATAR_TLE__
4+
#include "Conversion/TleToLLVM.h"
5+
#include "Dialect.h"
6+
#endif
37
#include "Dialect/TritonILUVATARGPU/IR/Dialect.h"
48
#include "PatternTritonGPUOpToLLVM.h"
59
#include "TargetInfo.h"
@@ -53,6 +57,9 @@ class TritonLLVMConversionTarget : public ConversionTarget {
5357
addIllegalDialect<triton::gpu::TritonGPUDialect>();
5458
addIllegalDialect<triton::nvidia_gpu::TritonNvidiaGPUDialect>();
5559
addIllegalDialect<mlir::gpu::GPUDialect>();
60+
#ifdef __ILUVATAR_TLE__
61+
mlir::triton::iluvatar_tle::addIllegalDialects(*this);
62+
#endif
5663
addLegalOp<mlir::UnrealizedConversionCastOp>();
5764
}
5865
};
@@ -69,6 +76,9 @@ struct ConvertTritonILUVATARGPUToLLVM
6976
registry
7077
.insert<LLVM::LLVMDialect, NVVM::NVVMDialect, mlir::ROCDL::ROCDLDialect,
7178
mlir::triton::iluvatargpu::TritonILUVATARGPUDialect>();
79+
#ifdef __ILUVATAR_TLE__
80+
mlir::triton::iluvatar_tle::registerDialects(registry);
81+
#endif
7282
}
7383

7484
void runOnOperation() override {
@@ -155,6 +165,10 @@ struct ConvertTritonILUVATARGPUToLLVM
155165
populateFunc(typeConverter, patterns, targetInfo, benefit);
156166
};
157167

168+
#ifdef __ILUVATAR_TLE__
169+
mlir::triton::iluvatar_tle::populateTleToLLVMPatterns(
170+
typeConverter, targetInfo, patterns, commonBenefit);
171+
#endif
158172
mlir::triton::populateConvertLayoutOpToLLVMPatterns(
159173
typeConverter, targetInfo, patterns, commonBenefit);
160174
ILUVATAR::populateDotOpToLLVMPatterns(typeConverter, patterns,

‎third_party/iluvatar/bin/RegisterTritonDialects.h‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,7 @@
11
#pragma once
2+
#ifdef __ILUVATAR_TLE__
3+
#include "Dialect.h"
4+
#endif
25
#include "amd/include/Dialect/TritonAMDGPU/IR/Dialect.h"
36
#include "amd/include/TritonAMDGPUTransforms/Passes.h"
47
#include "nvidia/include/Dialect/NVGPU/IR/Dialect.h"
@@ -62,7 +65,9 @@ void registerTestScopeIdAllocationPass();
6265
} // namespace mlir
6366

6467
inline void registerTritonDialects(mlir::DialectRegistry &registry) {
68+
#ifndef __ILUVATAR_TLE__
6569
mlir::registerAllPasses();
70+
#endif
6671
mlir::triton::registerTritonPasses();
6772
mlir::triton::gpu::registerTritonGPUPasses();
6873
mlir::triton::nvidia_gpu::registerTritonNvidiaGPUPasses();
@@ -148,4 +153,7 @@ inline void registerTritonDialects(mlir::DialectRegistry &registry) {
148153
mlir::triton::proton::ProtonDialect,
149154
mlir::triton::proton::gpu::ProtonGPUDialect, mlir::ROCDL::ROCDLDialect,
150155
mlir::triton::gluon::GluonDialect>();
156+
#ifdef __ILUVATAR_TLE__
157+
mlir::triton::iluvatar_tle::registerDialects(registry);
158+
#endif
151159
}

‎third_party/iluvatar/include/triton/Dialect/TritonGPU/Transforms/TritonGPUConversion.h‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -16,8 +16,15 @@ class TritonGPUTypeConverter : public TypeConverter {
1616
TritonGPUTypeConverter(MLIRContext *context, int numWarps, int threadsPerWarp,
1717
int numCTAs, bool enableSourceRemat);
1818
int getNumWarps() const { return numWarps; }
19+
#ifdef __ILUVATAR_TLE__
20+
int getNumWarps(Value value) const;
21+
#endif
1922
int getThreadsPerWarp() const { return threadsPerWarp; }
2023
int getNumCTAs() const { return numCTAs; }
24+
#ifdef __ILUVATAR_TLE__
25+
RankedTensorType convertRankedTensorType(RankedTensorType type,
26+
int contextualNumWarps) const;
27+
#endif
2128

2229
private:
2330
MLIRContext *context;

0 commit comments

Comments
 (0)