Skip to content

Commit 3f456fa

Browse files
whitneywhtsangwhitneywhtsang
committed
Merge commit '76ed94df1924b2262be9b37d778b6e0ccccb1180'
2 parents 7e59486 + 76ed94d commit 3f456fa

File tree

34 files changed

+220
-1092
lines changed

34 files changed

+220
-1092
lines changed

‎bin/RegisterTritonDialects.h‎

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -85,7 +85,6 @@ inline void registerTritonDialects(mlir::DialectRegistry &registry) {
8585
mlir::registerTritonAMDGPUAccelerateMatmul();
8686
mlir::registerTritonAMDGPUOptimizeEpilogue();
8787
mlir::registerTritonAMDGPUReorderInstructions();
88-
mlir::registerTritonAMDGPUStreamPipeline();
8988
mlir::registerTritonAMDGPUStreamPipelineV2();
9089
mlir::registerTritonAMDGPUCanonicalizePointers();
9190

‎include/triton/Conversion/TritonToTritonGPU/TritonToTritonGPUPass.h‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@
33

44
#include <memory>
55
#include <optional>
6+
#include <string>
67

78
namespace mlir {
89

‎include/triton/Dialect/Triton/IR/TritonOps.td‎

Lines changed: 2 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -460,17 +460,12 @@ def TT_ReshapeOp : TT_Op<"reshape", [Pure,
460460
If efficient_layout is set, this is a hint that the destination layout should be kept for performance reason.
461461
The compiler is still free to change it for better performance.
462462
}];
463-
let arguments = (ins TT_Tensor:$src, BoolAttr:$allow_reorder, OptionalAttr<UnitAttr>:$efficient_layout);
463+
let arguments = (ins TT_Tensor:$src, UnitAttr:$allow_reorder, UnitAttr:$efficient_layout);
464464
let results = (outs TT_Tensor:$result);
465-
let assemblyFormat = "$src attr-dict `:` type($src) `->` type($result)";
465+
let assemblyFormat = "$src (`allow_reorder` $allow_reorder^)? (`efficient_layout` $efficient_layout^)? attr-dict `:` type($src) `->` type($result)";
466466
let hasCanonicalizeMethod = 1;
467467
let hasFolder = 1;
468468
let hasVerifier = 1;
469-
let builders = [
470-
OpBuilder<(ins "Type":$type, "Value":$src, "bool":$allow_reorder),
471-
[{
472-
build($_builder, $_state, type, src, allow_reorder, /*efficient_layout=*/UnitAttr());
473-
}]>];
474469
}
475470

476471
def TT_BroadcastOp : TT_Op<"broadcast", [Pure,

‎lib/Dialect/Triton/IR/Ops.cpp‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -678,7 +678,7 @@ LogicalResult canonicalizeViewOrBroadcast(OpType op,
678678
}
679679

680680
LogicalResult ReshapeOp::canonicalize(ReshapeOp op, PatternRewriter &rewriter) {
681-
if (!op.getAllowReorder() || op.getEfficientLayout().has_value())
681+
if (!op.getAllowReorder() || op.getEfficientLayout())
682682
return failure();
683683
return canonicalizeViewOrBroadcast(op, rewriter);
684684
}

‎lib/Dialect/TritonGPU/IR/Dialect.cpp‎

Lines changed: 2 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -2764,7 +2764,7 @@ struct CanonicalizeConvertFromReshape
27642764
return failure();
27652765
if (isExpensiveView(convert.getSrc().getType(), op.getType()))
27662766
return failure();
2767-
if (!op.getAllowReorder() || op.getEfficientLayout().has_value())
2767+
if (!op.getAllowReorder() || op.getEfficientLayout())
27682768
return failure();
27692769

27702770
rewriter.replaceOpWithNewOp<triton::ReshapeOp>(
@@ -2885,8 +2885,7 @@ struct CanonicalizeConvertFromConvert
28852885

28862886
// cvt(reshape) -> reshape
28872887
if (auto reshape = dyn_cast<ReshapeOp>(arg)) {
2888-
if (!reshape.getAllowReorder() ||
2889-
reshape.getEfficientLayout().has_value() ||
2888+
if (!reshape.getAllowReorder() || reshape.getEfficientLayout() ||
28902889
isExpensiveView(reshape.getSrc().getType(), op.getType()))
28912890
return failure();
28922891

‎lib/Dialect/TritonGPU/Transforms/OptimizeThreadLocality.cpp‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -314,8 +314,8 @@ class TritonGPUOptimizeThreadLocalityPass
314314
IRMapping mapping;
315315
for (auto operand : reduce.getOperands()) {
316316
auto viewOp = builder.create<triton::ReshapeOp>(
317-
reduce.getLoc(), viewOpTensorType, operand, /*allowReorder=*/true);
318-
viewOp.setEfficientLayout(true);
317+
reduce.getLoc(), viewOpTensorType, operand,
318+
/*allowReorder=*/true, /*efficientLayout=*/true);
319319
mapping.map(operand, viewOp);
320320
}
321321

‎lib/Dialect/TritonGPU/Transforms/Utility.cpp‎

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -556,8 +556,7 @@ bool canFoldIntoConversion(Operation *op, Attribute targetEncoding) {
556556
RankedTensorType newDstType =
557557
RankedTensorType::get(reshapeDstType.getShape(),
558558
reshapeDstType.getElementType(), targetEncoding);
559-
return reshape.getAllowReorder() &&
560-
!reshape.getEfficientLayout().has_value() &&
559+
return reshape.getAllowReorder() && !reshape.getEfficientLayout() &&
561560
!triton::gpu::isExpensiveView(reshape.getSrc().getType(),
562561
newDstType);
563562
}

‎python/setup.py‎

Lines changed: 22 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -284,7 +284,8 @@ def download_and_copy(name, src_path, dst_path, variable, version, url_func):
284284
arch = {"x86_64": "64", "arm64": "aarch64", "aarch64": "aarch64"}[platform.machine()]
285285
except KeyError:
286286
arch = platform.machine()
287-
url = url_func(arch, version)
287+
supported = {"Linux": "linux", "Darwin": "linux"}
288+
url = url_func(supported[system], arch, version)
288289
tmp_path = os.path.join(triton_cache_path, "nvidia", name) # path to cache the download
289290
dst_path = os.path.join(base_dir, os.pardir, "third_party", "nvidia", "backend", dst_path) # final binary path
290291
platform_name = "sbsa-linux" if arch == "aarch64" else "x86_64-linux"
@@ -500,61 +501,62 @@ def get_platform_dependent_src_path(subdir):
500501

501502
download_and_copy(
502503
name="ptxas", src_path="bin/ptxas", dst_path="bin/ptxas", variable="TRITON_PTXAS_PATH",
503-
version=NVIDIA_TOOLCHAIN_VERSION["ptxas"], url_func=lambda arch, version:
504+
version=NVIDIA_TOOLCHAIN_VERSION["ptxas"], url_func=lambda system, arch, version:
504505
((lambda version_major, version_minor1, version_minor2:
505-
f"https://anaconda.org/nvidia/cuda-nvcc-tools/{version}/download/linux-{arch}/cuda-nvcc-tools-{version}-0.tar.bz2"
506+
f"https://anaconda.org/nvidia/cuda-nvcc-tools/{version}/download/{system}-{arch}/cuda-nvcc-tools-{version}-0.tar.bz2"
506507
if int(version_major) >= 12 and int(version_minor1) >= 5 else
507-
f"https://anaconda.org/nvidia/cuda-nvcc/{version}/download/linux-{arch}/cuda-nvcc-{version}-0.tar.bz2")
508+
f"https://anaconda.org/nvidia/cuda-nvcc/{version}/download/{system}-{arch}/cuda-nvcc-{version}-0.tar.bz2")
508509
(*version.split('.'))))
509510
download_and_copy(
510511
name="cuobjdump",
511512
src_path="bin/cuobjdump",
512513
dst_path="bin/cuobjdump",
513514
variable="TRITON_CUOBJDUMP_PATH",
514515
version=NVIDIA_TOOLCHAIN_VERSION["cuobjdump"],
515-
url_func=lambda arch, version:
516-
f"https://anaconda.org/nvidia/cuda-cuobjdump/{version}/download/linux-{arch}/cuda-cuobjdump-{version}-0.tar.bz2",
516+
url_func=lambda system, arch, version:
517+
f"https://anaconda.org/nvidia/cuda-cuobjdump/{version}/download/{system}-{arch}/cuda-cuobjdump-{version}-0.tar.bz2",
517518
)
518519
download_and_copy(
519520
name="nvdisasm",
520521
src_path="bin/nvdisasm",
521522
dst_path="bin/nvdisasm",
522523
variable="TRITON_NVDISASM_PATH",
523524
version=NVIDIA_TOOLCHAIN_VERSION["nvdisasm"],
524-
url_func=lambda arch, version:
525-
f"https://anaconda.org/nvidia/cuda-nvdisasm/{version}/download/linux-{arch}/cuda-nvdisasm-{version}-0.tar.bz2",
525+
url_func=lambda system, arch, version:
526+
f"https://anaconda.org/nvidia/cuda-nvdisasm/{version}/download/{system}-{arch}/cuda-nvdisasm-{version}-0.tar.bz2",
526527
)
527528
download_and_copy(
528529
name="cudacrt", src_path=get_platform_dependent_src_path("include"), dst_path="include",
529-
variable="TRITON_CUDACRT_PATH", version=NVIDIA_TOOLCHAIN_VERSION["cudacrt"], url_func=lambda arch, version:
530+
variable="TRITON_CUDACRT_PATH", version=NVIDIA_TOOLCHAIN_VERSION["cudacrt"], url_func=lambda system, arch, version:
530531
((lambda version_major, version_minor1, version_minor2:
531-
f"https://anaconda.org/nvidia/cuda-crt-dev_linux-{arch}/{version}/download/noarch/cuda-crt-dev_linux-{arch}-{version}-0.tar.bz2"
532+
f"https://anaconda.org/nvidia/cuda-crt-dev_{system}-{arch}/{version}/download/noarch/cuda-crt-dev_{system}-{arch}-{version}-0.tar.bz2"
532533
if int(version_major) >= 12 and int(version_minor1) >= 5 else
533-
f"https://anaconda.org/nvidia/cuda-nvcc/{version}/download/linux-{arch}/cuda-nvcc-{version}-0.tar.bz2")
534+
f"https://anaconda.org/nvidia/cuda-nvcc/{version}/download/{system}-{arch}/cuda-nvcc-{version}-0.tar.bz2")
534535
(*version.split('.'))))
535536
download_and_copy(
536537
name="cudart", src_path=get_platform_dependent_src_path("include"), dst_path="include",
537-
variable="TRITON_CUDART_PATH", version=NVIDIA_TOOLCHAIN_VERSION["cudart"], url_func=lambda arch, version:
538+
variable="TRITON_CUDART_PATH", version=NVIDIA_TOOLCHAIN_VERSION["cudart"], url_func=lambda system, arch, version:
538539
((lambda version_major, version_minor1, version_minor2:
539-
f"https://anaconda.org/nvidia/cuda-cudart-dev_linux-{arch}/{version}/download/noarch/cuda-cudart-dev_linux-{arch}-{version}-0.tar.bz2"
540+
f"https://anaconda.org/nvidia/cuda-cudart-dev_{system}-{arch}/{version}/download/noarch/cuda-cudart-dev_{system}-{arch}-{version}-0.tar.bz2"
540541
if int(version_major) >= 12 and int(version_minor1) >= 5 else
541-
f"https://anaconda.org/nvidia/cuda-cudart-dev/{version}/download/linux-{arch}/cuda-cudart-dev-{version}-0.tar.bz2"
542+
f"https://anaconda.org/nvidia/cuda-cudart-dev/{version}/download/{system}-{arch}/cuda-cudart-dev-{version}-0.tar.bz2"
542543
)(*version.split('.'))))
543544
download_and_copy(
544545
name="cupti", src_path=get_platform_dependent_src_path("include"), dst_path="include",
545-
variable="TRITON_CUPTI_INCLUDE_PATH", version=NVIDIA_TOOLCHAIN_VERSION["cupti"], url_func=lambda arch, version:
546+
variable="TRITON_CUPTI_INCLUDE_PATH", version=NVIDIA_TOOLCHAIN_VERSION["cupti"],
547+
url_func=lambda system, arch, version:
546548
((lambda version_major, version_minor1, version_minor2:
547-
f"https://anaconda.org/nvidia/cuda-cupti-dev/{version}/download/linux-{arch}/cuda-cupti-dev-{version}-0.tar.bz2"
549+
f"https://anaconda.org/nvidia/cuda-cupti-dev/{version}/download/{system}-{arch}/cuda-cupti-dev-{version}-0.tar.bz2"
548550
if int(version_major) >= 12 and int(version_minor1) >= 5 else
549-
f"https://anaconda.org/nvidia/cuda-cupti/{version}/download/linux-{arch}/cuda-cupti-{version}-0.tar.bz2")
551+
f"https://anaconda.org/nvidia/cuda-cupti/{version}/download/{system}-{arch}/cuda-cupti-{version}-0.tar.bz2")
550552
(*version.split('.'))))
551553
download_and_copy(
552554
name="cupti", src_path=get_platform_dependent_src_path("lib"), dst_path="lib/cupti",
553-
variable="TRITON_CUPTI_LIB_PATH", version=NVIDIA_TOOLCHAIN_VERSION["cupti"], url_func=lambda arch, version:
555+
variable="TRITON_CUPTI_LIB_PATH", version=NVIDIA_TOOLCHAIN_VERSION["cupti"], url_func=lambda system, arch, version:
554556
((lambda version_major, version_minor1, version_minor2:
555-
f"https://anaconda.org/nvidia/cuda-cupti-dev/{version}/download/linux-{arch}/cuda-cupti-dev-{version}-0.tar.bz2"
557+
f"https://anaconda.org/nvidia/cuda-cupti-dev/{version}/download/{system}-{arch}/cuda-cupti-dev-{version}-0.tar.bz2"
556558
if int(version_major) >= 12 and int(version_minor1) >= 5 else
557-
f"https://anaconda.org/nvidia/cuda-cupti/{version}/download/linux-{arch}/cuda-cupti-{version}-0.tar.bz2")
559+
f"https://anaconda.org/nvidia/cuda-cupti/{version}/download/{system}-{arch}/cuda-cupti-{version}-0.tar.bz2")
558560
(*version.split('.'))))
559561

560562
backends = [*BackendInstaller.copy(["intel", "nvidia", "amd"]), *BackendInstaller.copy_externals()]

‎python/test/unit/language/test_compile_errors.py‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -154,7 +154,8 @@ def kernel():
154154
try:
155155
inner = e.value.__cause__
156156
outer = e.value
157-
assert "/core.py" in '\n'.join(traceback.format_tb(inner.__traceback__)), "error should point inside core.py"
157+
assert f"{os.sep}core.py" in '\n'.join(traceback.format_tb(
158+
inner.__traceback__)), "error should point inside core.py"
158159

159160
assert "at 2:4:" in str(outer), "error should point to expand_dims call"
160161
assert "<source unavailable>" not in str(outer)

‎python/test/unit/language/test_subprocess.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -53,7 +53,7 @@ def test_print(func_type: str, data_type: str, device: str):
5353
assert proc.stderr == b''
5454
return
5555

56-
outs = [line for line in proc.stdout.decode("UTF-8").split("\n") if line]
56+
outs = [line for line in proc.stdout.decode("UTF-8").splitlines() if line]
5757
# The total number of elements in the 1-D tensor to print.
5858
N = 128
5959

0 commit comments

Comments
 (0)