From 80c2b4412d8e042ef30015bd14a6580dd2820385 Mon Sep 17 00:00:00 2001 From: Stardep <1486216685@qq.com> Date: Fri, 14 Aug 2026 01:21:58 +0800 Subject: [PATCH 1/4] [SPEC][MTHREADS] Unify C++ backend specialization --- CMakeLists.txt | 5 + cmake/FlagTreeBackendSpec.cmake | 23 +- cmake/FlagTreeOptions.cmake | 23 +- .../Dialect/Gluon/Transforms/CMakeLists.txt | 2 +- .../triton/Dialect/Triton/IR/CMakeLists.txt | 4 +- .../Dialect/TritonGPU/IR/CMakeLists.txt | 4 +- .../TritonInstrument/IR/CMakeLists.txt | 2 +- .../Dialect/TritonNvidiaGPU/IR/CMakeLists.txt | 6 +- lib/Dialect/Triton/Transforms/CMakeLists.txt | 2 +- third_party/mthreads/CMakeLists.txt | 59 +- third_party/mthreads/include/CMakeLists.txt | 1 - .../include/triton/Analysis/AxisInfo.h | 271 ---- .../mthreads/include/triton/CMakeLists.txt | 3 - .../include/triton/Conversion/CMakeLists.txt | 2 - .../include/triton/Conversion/MLIRTypes.h | 46 - .../AllocateSharedMemoryUtility.h | 17 - .../Conversion/TritonGPUToLLVM/AsmFormat.h | 27 - .../Conversion/TritonGPUToLLVM/CMakeLists.txt | 3 - .../TritonGPUToLLVM/FMADotUtility.h | 35 - .../Conversion/TritonGPUToLLVM/Passes.h | 25 - .../Conversion/TritonGPUToLLVM/Passes.td | 45 - .../TritonGPUToLLVM/TypeConverter.h | 39 - .../TritonToTritonGPU/CMakeLists.txt | 3 - .../Conversion/TritonToTritonGPU/Passes.h | 15 - .../Conversion/TritonToTritonGPU/Passes.td | 56 - .../include/triton/Dialect/CMakeLists.txt | 7 - .../triton/Dialect/Gluon/CMakeCache.txt | 2 - .../triton/Dialect/Gluon/CMakeLists.txt | 2 - .../triton/Dialect/Gluon/IR/CMakeLists.txt | 17 - .../include/triton/Dialect/Gluon/IR/Dialect.h | 11 - .../triton/Dialect/Gluon/IR/GluonAttrDefs.td | 23 - .../triton/Dialect/Gluon/IR/GluonDialect.td | 22 - .../triton/Dialect/Gluon/IR/GluonOps.td | 32 - .../Dialect/Gluon/Transforms/CMakeLists.txt | 3 - .../Gluon/Transforms/InferLayoutUtils.h | 20 - .../triton/Dialect/Gluon/Transforms/Passes.h | 13 - .../triton/Dialect/NVGPU/CMakeLists.txt | 1 - .../triton/Dialect/NVWS/CMakeLists.txt | 2 - .../triton/Dialect/Triton/CMakeLists.txt | 2 - .../triton/Dialect/Triton/IR/CMakeLists.txt | 31 - .../Dialect/Triton/IR/DiscardableAttributes.h | 15 - .../triton/Dialect/Triton/IR/Interfaces.h | 45 - .../triton/Dialect/Triton/IR/OpInterfaces.h | 24 - .../Dialect/Triton/IR/TritonAttrDefs.td | 154 -- .../triton/Dialect/Triton/IR/TritonDialect.td | 60 - .../Dialect/Triton/IR/TritonOpInterfaces.td | 118 -- .../Triton/Transforms/ArithTypeConversion.h | 18 - .../Dialect/Triton/Transforms/CMakeLists.txt | 3 - .../Transforms/FunctionTypeConversion.h | 19 - .../Dialect/Triton/Transforms/LoopPeeling.h | 18 - .../triton/Dialect/Triton/Transforms/Passes.h | 19 - .../Dialect/Triton/Transforms/Passes.td | 93 -- .../triton/Dialect/TritonGPU/CMakeLists.txt | 2 - .../Dialect/TritonGPU/IR/CMakeLists.txt | 45 - .../Dialect/TritonGPU/IR/TritonGPUDialect.td | 41 - .../TritonGPU/IR/TritonGPUOpInterfaces.td | 29 - .../TritonGPU/IR/TritonGPUTypeInterfaces.td | 23 - .../triton/Dialect/TritonGPU/IR/Types.h | 14 - .../TritonGPU/Transforms/CMakeLists.txt | 3 - .../Transforms/DecomposeScaledBlocked.h | 47 - .../Transforms/LayoutPropagationUtility.h | 21 - .../Transforms/MMAv5PipelineUtility.h | 83 - .../Dialect/TritonGPU/Transforms/Partition.h | 127 -- .../TritonGPU/Transforms/PartitionBuilder.h | 49 - .../TritonGPU/Transforms/PipelineExpander.h | 111 -- .../TritonGPU/Transforms/PipeliningUtility.h | 189 --- .../TritonGPU/Transforms/WarpSpecialization.h | 24 - .../Dialect/TritonInstrument/CMakeLists.txt | 2 - .../TritonInstrument/IR/CMakeLists.txt | 15 - .../Dialect/TritonInstrument/IR/Dialect.h | 14 - .../TritonInstrument/IR/TritonInstrument.md | 86 - .../IR/TritonInstrumentAttrDefs.td | 15 - .../IR/TritonInstrumentDialect.td | 11 - .../Transforms/CMakeLists.txt | 3 - .../TritonInstrument/Transforms/Passes.h | 22 - .../TritonInstrument/Transforms/Passes.td | 16 - .../Dialect/TritonNvidiaGPU/CMakeLists.txt | 2 - .../Dialect/TritonNvidiaGPU/IR/CMakeLists.txt | 27 - .../TritonNvidiaGPU/IR/TensorMemoryUtils.h | 37 - .../TritonNvidiaGPU/Transforms/CMakeLists.txt | 3 - .../TritonNvidiaGPU/Transforms/Passes.h | 46 - .../TritonNvidiaGPU/Transforms/Passes.td | 187 --- .../include/triton/Target/CMakeLists.txt | 1 - .../triton/Target/LLVMIR/CMakeLists.txt | 3 - .../include/triton/Target/LLVMIR/Passes.h | 18 - .../include/triton/Target/LLVMIR/Passes.td | 21 - .../include/triton/Tools/GenericSwizzling.h | 56 - .../include/triton/Tools/LayoutUtils.h | 190 --- .../mthreads/include/triton/Tools/StrUtil.h | 54 - .../mthreads/lib/Analysis/CMakeLists.txt | 23 - third_party/mthreads/lib/CMakeLists.txt | 5 - .../mthreads/lib/Conversion/CMakeLists.txt | 3 - .../TritonGPUToLLVM/AllocateSharedMemory.cpp | 27 - .../AllocateSharedMemoryUtility.cpp | 34 - .../Conversion/TritonGPUToLLVM/CMakeLists.txt | 41 - .../TritonGPUToLLVM/ControlFlowOpToLLVM.cpp | 171 -- .../DotOpToLLVM/FMADotUtility.cpp | 170 -- .../TritonGPUToLLVM/MakeRangeOpToLLVM.cpp | 54 - .../TritonGPUToLLVM/PrintOpToLLVM.cpp | 243 --- .../TritonGPUToLLVM/ReduceScanCommon.h | 163 -- .../TritonGPUToLLVM/SPMDOpToLLVM.cpp | 37 - .../TritonGPUToLLVM/TypeConverter.cpp | 77 - .../TritonInstrumentToLLVM/CMakeLists.txt | 12 - .../TritonToTritonGPU/CMakeLists.txt | 26 - .../mthreads/lib/Dialect/CMakeLists.txt | 7 - .../mthreads/lib/Dialect/Gluon/CMakeLists.txt | 2 - .../lib/Dialect/Gluon/IR/CMakeLists.txt | 10 - .../mthreads/lib/Dialect/Gluon/IR/Dialect.cpp | 138 -- .../Dialect/Gluon/Transforms/CMakeLists.txt | 17 - .../Gluon/Transforms/InferLayoutUtils.cpp | 251 --- .../lib/Dialect/Gluon/Transforms/Inline.cpp | 29 - .../Gluon/Transforms/ResolveAutoEncodings.cpp | 71 - .../Gluon/Transforms/SimplifyControlFlow.cpp | 49 - .../mthreads/lib/Dialect/NVGPU/CMakeLists.txt | 1 - .../mthreads/lib/Dialect/NVWS/CMakeLists.txt | 2 - .../lib/Dialect/Triton/CMakeLists.txt | 2 - .../lib/Dialect/Triton/IR/CMakeLists.txt | 23 - .../lib/Dialect/Triton/IR/Canonicalize.td | 17 - .../lib/Dialect/Triton/IR/Dialect.cpp | 77 - .../Triton/IR/DiscardableAttributes.cpp | 17 - .../lib/Dialect/Triton/IR/OpInterfaces.cpp | 77 - .../Triton/Transforms/ArithTypeConversion.cpp | 51 - .../Dialect/Triton/Transforms/CMakeLists.txt | 27 - .../Transforms/FunctionTypeConversion.cpp | 163 -- .../Triton/Transforms/LoopAwareCSE.cpp | 178 --- .../Transforms/LoopInvariantCodeMotion.cpp | 82 - .../Dialect/Triton/Transforms/LoopPeeling.cpp | 67 - .../Triton/Transforms/ReorderBroadcast.cpp | 230 --- .../lib/Dialect/TritonGPU/CMakeLists.txt | 2 - .../lib/Dialect/TritonGPU/IR/CMakeLists.txt | 18 - .../TritonGPU/Transforms/CMakeLists.txt | 56 - .../Transforms/CombineTensorSelectAndIf.cpp | 176 --- .../Transforms/DecomposeScaledBlocked.cpp | 261 --- .../TritonGPU/Transforms/FuseNestedLoops.cpp | 1222 -------------- .../TritonGPU/Transforms/HoistTMEMAlloc.cpp | 586 ------- .../Transforms/LayoutPropagationUtility.cpp | 49 - .../Pipeliner/TestPipelineLowerLoop.cpp | 32 - .../Transforms/ReorderInstructions.cpp | 178 --- .../WarpSpecialization/PartitionBuilder.cpp | 36 - .../Dialect/TritonInstrument/CMakeLists.txt | 2 - .../TritonInstrument/IR/CMakeLists.txt | 14 - .../Dialect/TritonInstrument/IR/Dialect.cpp | 17 - .../lib/Dialect/TritonInstrument/IR/Ops.cpp | 8 - .../Transforms/CMakeLists.txt | 16 - .../Dialect/TritonNvidiaGPU/CMakeLists.txt | 2 - .../Dialect/TritonNvidiaGPU/IR/CMakeLists.txt | 15 - .../TritonNvidiaGPU/Transforms/CMakeLists.txt | 25 - .../Transforms/CheckMatmulTwoCTAs.cpp | 63 - .../mthreads/lib/Target/CMakeLists.txt | 1 - .../mthreads/lib/Target/LLVMIR/CMakeLists.txt | 31 - .../Target/LLVMIR/LLVMIRBreakPhiStruct.cpp | 60 - .../mthreads/lib/Target/LLVMIR/LLVMPasses.h | 16 - third_party/mthreads/lib/Tools/CMakeLists.txt | 13 - .../mthreads/lib/Tools/GenericSwizzling.cpp | 713 --------- .../mthreads/lib/Tools/LayoutUtils.cpp | 582 ------- .../mthreads/lib/Tools/LinearLayout.cpp | 1407 ----------------- .../include/triton/Analysis/Alias.h | 0 .../include/triton/Analysis/Allocation.h | 0 .../include/triton/Analysis/BufferRegion.h | 0 .../include/triton/Analysis/Membar.h | 0 .../include/triton/Analysis/Utility.h | 0 .../TritonGPUToLLVM/ElementwiseOpToLLVMBase.h | 0 .../PatternTritonGPUOpToLLVM.h | 0 .../TritonGPUToLLVM/TargetInfoBase.h | 0 .../Conversion/TritonGPUToLLVM/Utility.h | 0 .../TritonGPUToLLVM/WarpSpecializeUtility.h | 0 .../triton/Dialect/Gluon/Transforms/Passes.td | 0 .../triton/Dialect/NVGPU/IR/CMakeLists.txt | 0 .../include/triton/Dialect/NVGPU/IR/Dialect.h | 0 .../triton/Dialect/NVGPU/IR/NVGPUAttrDefs.td | 0 .../triton/Dialect/NVGPU/IR/NVGPUDialect.td | 0 .../triton/Dialect/NVGPU/IR/NVGPUOps.td | 0 .../triton/Dialect/NVWS/IR/CMakeLists.txt | 0 .../include/triton/Dialect/NVWS/IR/Dialect.h | 0 .../triton/Dialect/NVWS/IR/NVWSAttrDefs.td | 0 .../triton/Dialect/NVWS/IR/NVWSDialect.td | 0 .../Dialect/NVWS/IR/NVWSOpInterfaces.td | 0 .../include/triton/Dialect/NVWS/IR/NVWSOps.td | 0 .../triton/Dialect/NVWS/IR/NVWSTypes.td | 0 .../Dialect/NVWS/Transforms/CMakeLists.txt | 0 .../triton/Dialect/NVWS/Transforms/Passes.h | 0 .../triton/Dialect/NVWS/Transforms/Passes.td | 0 .../triton/Dialect/Triton/IR/CMakeLists.txt | 4 + .../triton/Dialect/Triton/IR/Dialect.h | 0 .../include/triton/Dialect/Triton/IR/Traits.h | 0 .../Dialect/Triton/IR/TritonInterfaces.td | 0 .../triton/Dialect/Triton/IR/TritonOps.td | 0 .../Dialect/Triton/IR/TritonTypeInterfaces.td | 0 .../triton/Dialect/Triton/IR/TritonTypes.td | 0 .../include/triton/Dialect/Triton/IR/Types.h | 0 .../triton/Dialect/Triton/IR/Utility.h | 0 .../triton/Dialect/TritonGPU/IR/Attributes.h | 0 .../Dialect/TritonGPU/IR/CGAEncodingAttr.h | 0 .../Dialect/TritonGPU/IR/CGAEncodingAttr.td | 0 .../Dialect/TritonGPU/IR/CMakeLists.txt | 8 + .../triton/Dialect/TritonGPU/IR/Dialect.h | 0 .../TritonGPU/IR/LinearLayoutConversions.h | 0 .../triton/Dialect/TritonGPU/IR/Traits.h | 0 .../Dialect/TritonGPU/IR/TritonGPUAttrBase.td | 0 .../Dialect/TritonGPU/IR/TritonGPUAttrDefs.td | 0 .../TritonGPU/IR/TritonGPUAttrImpls.td | 0 .../Dialect/TritonGPU/IR/TritonGPUEnums.td | 0 .../TritonGPU/IR/TritonGPUInterfaces.h | 0 .../Dialect/TritonGPU/IR/TritonGPUOps.td | 0 .../Dialect/TritonGPU/IR/TritonGPUTypes.td | 0 .../TritonGPU/Transforms/CoalesceUtils.h | 0 .../Transforms/PartitionSchedulingUtility.h | 0 .../Dialect/TritonGPU/Transforms/Passes.h | 0 .../Dialect/TritonGPU/Transforms/Passes.td | 0 .../Dialect/TritonGPU/Transforms/Schedule.h | 0 .../Transforms/TritonGPUConversion.h | 0 .../Dialect/TritonGPU/Transforms/Utility.h | 0 .../TritonInstrument/IR/FunctionBuilder.h | 0 .../IR/TritonInstrumentOps.td | 0 .../Dialect/TritonInstrument/IR/Utility.h | 0 .../Dialect/TritonNvidiaGPU/IR/CMakeLists.txt | 4 + .../Dialect/TritonNvidiaGPU/IR/Dialect.h | 0 .../IR/TritonNvidiaGPUAttrDefs.td | 0 .../IR/TritonNvidiaGPUDialect.td | 0 .../IR/TritonNvidiaGPUOpInterfaces.td | 0 .../TritonNvidiaGPU/IR/TritonNvidiaGPUOps.td | 0 .../IR/TritonNvidiaGPUTypes.td | 0 .../TritonNvidiaGPU/Transforms/TMAUtilities.h | 0 .../include/triton/Tools/LinearLayout.h | 0 .../include/triton/Tools/PluginUtils.h | 0 .../include/triton/Tools/Sys/GetEnv.hpp | 0 .../{ => spec_cpp}/lib/Analysis/Alias.cpp | 0 .../lib/Analysis/Allocation.cpp | 0 .../{ => spec_cpp}/lib/Analysis/AxisInfo.cpp | 0 .../lib/Analysis/BufferRegion.cpp | 0 .../{ => spec_cpp}/lib/Analysis/Membar.cpp | 0 .../{ => spec_cpp}/lib/Analysis/Utility.cpp | 0 .../TritonGPUToLLVM/AllocateWarpGroups.cpp | 0 .../TritonGPUToLLVM/AssertOpToLLVM.cpp | 0 .../TritonGPUToLLVM/ConvertLayoutOpToLLVM.cpp | 0 .../TritonGPUToLLVM/DotOpToLLVM/FMA.cpp | 0 .../TritonGPUToLLVM/ElementwiseOpToLLVM.cpp | 0 .../TritonGPUToLLVM/FuncOpToLLVM.cpp | 0 .../TritonGPUToLLVM/GatherOpToLLVM.cpp | 0 .../GlobalScratchMemoryAllocation.cpp | 0 .../TritonGPUToLLVM/HistogramOpToLLVM.cpp | 0 .../TritonGPUToLLVM/MemoryOpToLLVM.cpp | 0 .../TritonGPUToLLVM/ReduceOpToLLVM.cpp | 0 .../TritonGPUToLLVM/ScanOpToLLVM.cpp | 0 .../Conversion/TritonGPUToLLVM/Utility.cpp | 0 .../TritonGPUToLLVM/ViewOpToLLVM.cpp | 0 .../TritonGPUToLLVM/WarpSpecializeUtility.cpp | 0 .../InstrumentationToLLVM.cpp | 0 .../TritonToTritonGPU/RelayoutTritonGPU.cpp | 0 .../TritonToTritonGPU/TritonGPUConversion.cpp | 0 .../TritonToTritonGPUPass.cpp | 0 .../Dialect/Gluon/Transforms/Canonicalize.cpp | 0 .../Transforms/InferCoalescedEncodings.cpp | 0 .../lib/Dialect/NVGPU/IR/CMakeLists.txt | 0 .../lib/Dialect/NVGPU/IR/Dialect.cpp | 0 .../lib/Dialect/NVWS/IR/CMakeLists.txt | 0 .../lib/Dialect/NVWS/IR/Dialect.cpp | 0 .../lib/Dialect/NVWS/IR/Ops.cpp | 0 .../NVWS/Transforms/AssignStagePhase.cpp | 0 .../Dialect/NVWS/Transforms/CMakeLists.txt | 0 .../NVWS/Transforms/HoistTmemStore.cpp | 0 .../Dialect/NVWS/Transforms/InsertAref.cpp | 0 .../NVWS/Transforms/InsertTmemAref.cpp | 0 .../lib/Dialect/NVWS/Transforms/LowerAref.cpp | 0 .../NVWS/Transforms/LowerWarpGroup.cpp | 0 .../lib/Dialect/NVWS/Transforms/Utilities.cpp | 0 .../lib/Dialect/NVWS/Transforms/Utilities.h | 0 .../lib/Dialect/Triton/IR/Ops.cpp | 0 .../lib/Dialect/Triton/IR/Traits.cpp | 0 .../lib/Dialect/Triton/IR/Types.cpp | 0 .../lib/Dialect/Triton/IR/Utility.cpp | 0 .../lib/Dialect/Triton/Transforms/Combine.cpp | 0 .../lib/Dialect/Triton/Transforms/Combine.td | 0 .../Dialect/Triton/Transforms/LoopUnroll.cpp | 0 .../RewriteTensorDescriptorToPointer.cpp | 0 .../Transforms/RewriteTensorPointer.cpp | 0 .../lib/Dialect/TritonGPU/IR/Dialect.cpp | 0 .../TritonGPU/IR/LinearLayoutConversions.cpp | 0 .../lib/Dialect/TritonGPU/IR/Ops.cpp | 0 .../lib/Dialect/TritonGPU/IR/Types.cpp | 0 .../TritonGPU/Transforms/AccelerateMatmul.cpp | 0 .../Dialect/TritonGPU/Transforms/Coalesce.cpp | 0 .../Transforms/CoalesceAsyncCopy.cpp | 0 .../TritonGPU/Transforms/CoalesceUtils.cpp | 0 .../Dialect/TritonGPU/Transforms/F32DotTC.cpp | 0 .../Transforms/OptimizeAccumulatorInit.cpp | 0 .../Transforms/OptimizeDotOperands.cpp | 0 .../Transforms/OptimizeThreadLocality.cpp | 0 .../Transforms/Pipeliner/AssignLatencies.cpp | 0 .../Transforms/Pipeliner/LowerLoops.cpp | 0 .../Pipeliner/MMAv5PipelineUtility.cpp | 0 .../Transforms/Pipeliner/PipelineExpander.cpp | 0 .../Pipeliner/PipeliningUtility.cpp | 0 .../Transforms/Pipeliner/Schedule.cpp | 0 .../Transforms/Pipeliner/ScheduleLoops.cpp | 0 .../Pipeliner/SoftwarePipeliner.cpp | 0 .../Pipeliner/TMAStoresPipeline.cpp | 0 .../Transforms/Pipeliner/WGMMAPipeline.cpp | 0 .../Dialect/TritonGPU/Transforms/Prefetch.cpp | 0 .../Transforms/ReduceDataDuplication.cpp | 0 .../Transforms/RemoveLayoutConversions.cpp | 0 .../Dialect/TritonGPU/Transforms/Utility.cpp | 0 .../AutomaticWarpSpecialization.cpp | 0 .../OptimizePartitionWarps.cpp | 0 .../WarpSpecialization/Partition.cpp | 0 .../WarpSpecialization/PartitionLoops.cpp | 0 .../PartitionScheduling.cpp | 0 .../PartitionSchedulingUtility.cpp | 0 .../TritonInstrument/IR/FunctionBuilder.cpp | 0 .../Dialect/TritonInstrument/IR/Utility.cpp | 0 .../Transforms/ConcurrencySanitizer.cpp | 0 .../Dialect/TritonNvidiaGPU/IR/Dialect.cpp | 0 .../lib/Dialect/TritonNvidiaGPU/IR/Ops.cpp | 0 .../TritonNvidiaGPU/IR/TensorMemoryUtils.cpp | 0 .../Transforms/FenceInsertion.cpp | 0 .../Transforms/InterleaveTMem.cpp | 0 .../Transforms/MMALowering.cpp | 0 .../Transforms/OptimizeDescriptorEncoding.cpp | 0 .../Transforms/OptimizeTMemLayouts.cpp | 0 .../TritonNvidiaGPU/Transforms/PlanCTA.cpp | 0 .../Transforms/PromoteLHSToTMem.cpp | 0 .../Transforms/ProxyFenceInsertion.cpp | 0 .../Transforms/RemoveTMEMTokens.cpp | 0 .../Transforms/TMALowering.cpp | 0 .../Transforms/TMAUtilities.cpp | 0 .../Transforms/TensorMemoryAllocation.cpp | 0 .../lib/Target/LLVMIR/LLVMDILocalVariable.cpp | 0 .../lib/Target/LLVMIR/LLVMDIScope.cpp | 0 .../lib/Target/LLVMIR/LLVMDIUtils.cpp | 0 .../lib/Target/LLVMIR/LLVMDIUtils.h | 0 .../{ => spec_cpp}/lib/Tools/PluginUtils.cpp | 0 331 files changed, 127 insertions(+), 11293 deletions(-) delete mode 100644 third_party/mthreads/include/CMakeLists.txt delete mode 100644 third_party/mthreads/include/triton/Analysis/AxisInfo.h delete mode 100644 third_party/mthreads/include/triton/CMakeLists.txt delete mode 100644 third_party/mthreads/include/triton/Conversion/CMakeLists.txt delete mode 100644 third_party/mthreads/include/triton/Conversion/MLIRTypes.h delete mode 100644 third_party/mthreads/include/triton/Conversion/TritonGPUToLLVM/AllocateSharedMemoryUtility.h delete mode 100644 third_party/mthreads/include/triton/Conversion/TritonGPUToLLVM/AsmFormat.h delete mode 100644 third_party/mthreads/include/triton/Conversion/TritonGPUToLLVM/CMakeLists.txt delete mode 100644 third_party/mthreads/include/triton/Conversion/TritonGPUToLLVM/FMADotUtility.h delete mode 100644 third_party/mthreads/include/triton/Conversion/TritonGPUToLLVM/Passes.h delete mode 100644 third_party/mthreads/include/triton/Conversion/TritonGPUToLLVM/Passes.td delete mode 100644 third_party/mthreads/include/triton/Conversion/TritonGPUToLLVM/TypeConverter.h delete mode 100644 third_party/mthreads/include/triton/Conversion/TritonToTritonGPU/CMakeLists.txt delete mode 100644 third_party/mthreads/include/triton/Conversion/TritonToTritonGPU/Passes.h delete mode 100644 third_party/mthreads/include/triton/Conversion/TritonToTritonGPU/Passes.td delete mode 100644 third_party/mthreads/include/triton/Dialect/CMakeLists.txt delete mode 100644 third_party/mthreads/include/triton/Dialect/Gluon/CMakeCache.txt delete mode 100644 third_party/mthreads/include/triton/Dialect/Gluon/CMakeLists.txt delete mode 100644 third_party/mthreads/include/triton/Dialect/Gluon/IR/CMakeLists.txt delete mode 100644 third_party/mthreads/include/triton/Dialect/Gluon/IR/Dialect.h delete mode 100644 third_party/mthreads/include/triton/Dialect/Gluon/IR/GluonAttrDefs.td delete mode 100644 third_party/mthreads/include/triton/Dialect/Gluon/IR/GluonDialect.td delete mode 100644 third_party/mthreads/include/triton/Dialect/Gluon/IR/GluonOps.td delete mode 100644 third_party/mthreads/include/triton/Dialect/Gluon/Transforms/CMakeLists.txt delete mode 100644 third_party/mthreads/include/triton/Dialect/Gluon/Transforms/InferLayoutUtils.h delete mode 100644 third_party/mthreads/include/triton/Dialect/Gluon/Transforms/Passes.h delete mode 100644 third_party/mthreads/include/triton/Dialect/NVGPU/CMakeLists.txt delete mode 100644 third_party/mthreads/include/triton/Dialect/NVWS/CMakeLists.txt delete mode 100644 third_party/mthreads/include/triton/Dialect/Triton/CMakeLists.txt delete mode 100644 third_party/mthreads/include/triton/Dialect/Triton/IR/CMakeLists.txt delete mode 100644 third_party/mthreads/include/triton/Dialect/Triton/IR/DiscardableAttributes.h delete mode 100644 third_party/mthreads/include/triton/Dialect/Triton/IR/Interfaces.h delete mode 100644 third_party/mthreads/include/triton/Dialect/Triton/IR/OpInterfaces.h delete mode 100644 third_party/mthreads/include/triton/Dialect/Triton/IR/TritonAttrDefs.td delete mode 100644 third_party/mthreads/include/triton/Dialect/Triton/IR/TritonDialect.td delete mode 100644 third_party/mthreads/include/triton/Dialect/Triton/IR/TritonOpInterfaces.td delete mode 100644 third_party/mthreads/include/triton/Dialect/Triton/Transforms/ArithTypeConversion.h delete mode 100644 third_party/mthreads/include/triton/Dialect/Triton/Transforms/CMakeLists.txt delete mode 100644 third_party/mthreads/include/triton/Dialect/Triton/Transforms/FunctionTypeConversion.h delete mode 100644 third_party/mthreads/include/triton/Dialect/Triton/Transforms/LoopPeeling.h delete mode 100644 third_party/mthreads/include/triton/Dialect/Triton/Transforms/Passes.h delete mode 100644 third_party/mthreads/include/triton/Dialect/Triton/Transforms/Passes.td delete mode 100644 third_party/mthreads/include/triton/Dialect/TritonGPU/CMakeLists.txt delete mode 100644 third_party/mthreads/include/triton/Dialect/TritonGPU/IR/CMakeLists.txt delete mode 100644 third_party/mthreads/include/triton/Dialect/TritonGPU/IR/TritonGPUDialect.td delete mode 100644 third_party/mthreads/include/triton/Dialect/TritonGPU/IR/TritonGPUOpInterfaces.td delete mode 100644 third_party/mthreads/include/triton/Dialect/TritonGPU/IR/TritonGPUTypeInterfaces.td delete mode 100644 third_party/mthreads/include/triton/Dialect/TritonGPU/IR/Types.h delete mode 100644 third_party/mthreads/include/triton/Dialect/TritonGPU/Transforms/CMakeLists.txt delete mode 100644 third_party/mthreads/include/triton/Dialect/TritonGPU/Transforms/DecomposeScaledBlocked.h delete mode 100644 third_party/mthreads/include/triton/Dialect/TritonGPU/Transforms/LayoutPropagationUtility.h delete mode 100644 third_party/mthreads/include/triton/Dialect/TritonGPU/Transforms/MMAv5PipelineUtility.h delete mode 100644 third_party/mthreads/include/triton/Dialect/TritonGPU/Transforms/Partition.h delete mode 100644 third_party/mthreads/include/triton/Dialect/TritonGPU/Transforms/PartitionBuilder.h delete mode 100644 third_party/mthreads/include/triton/Dialect/TritonGPU/Transforms/PipelineExpander.h delete mode 100644 third_party/mthreads/include/triton/Dialect/TritonGPU/Transforms/PipeliningUtility.h delete mode 100644 third_party/mthreads/include/triton/Dialect/TritonGPU/Transforms/WarpSpecialization.h delete mode 100644 third_party/mthreads/include/triton/Dialect/TritonInstrument/CMakeLists.txt delete mode 100644 third_party/mthreads/include/triton/Dialect/TritonInstrument/IR/CMakeLists.txt delete mode 100644 third_party/mthreads/include/triton/Dialect/TritonInstrument/IR/Dialect.h delete mode 100644 third_party/mthreads/include/triton/Dialect/TritonInstrument/IR/TritonInstrument.md delete mode 100644 third_party/mthreads/include/triton/Dialect/TritonInstrument/IR/TritonInstrumentAttrDefs.td delete mode 100644 third_party/mthreads/include/triton/Dialect/TritonInstrument/IR/TritonInstrumentDialect.td delete mode 100644 third_party/mthreads/include/triton/Dialect/TritonInstrument/Transforms/CMakeLists.txt delete mode 100644 third_party/mthreads/include/triton/Dialect/TritonInstrument/Transforms/Passes.h delete mode 100644 third_party/mthreads/include/triton/Dialect/TritonInstrument/Transforms/Passes.td delete mode 100644 third_party/mthreads/include/triton/Dialect/TritonNvidiaGPU/CMakeLists.txt delete mode 100644 third_party/mthreads/include/triton/Dialect/TritonNvidiaGPU/IR/CMakeLists.txt delete mode 100644 third_party/mthreads/include/triton/Dialect/TritonNvidiaGPU/IR/TensorMemoryUtils.h delete mode 100644 third_party/mthreads/include/triton/Dialect/TritonNvidiaGPU/Transforms/CMakeLists.txt delete mode 100644 third_party/mthreads/include/triton/Dialect/TritonNvidiaGPU/Transforms/Passes.h delete mode 100644 third_party/mthreads/include/triton/Dialect/TritonNvidiaGPU/Transforms/Passes.td delete mode 100644 third_party/mthreads/include/triton/Target/CMakeLists.txt delete mode 100644 third_party/mthreads/include/triton/Target/LLVMIR/CMakeLists.txt delete mode 100644 third_party/mthreads/include/triton/Target/LLVMIR/Passes.h delete mode 100644 third_party/mthreads/include/triton/Target/LLVMIR/Passes.td delete mode 100644 third_party/mthreads/include/triton/Tools/GenericSwizzling.h delete mode 100644 third_party/mthreads/include/triton/Tools/LayoutUtils.h delete mode 100644 third_party/mthreads/include/triton/Tools/StrUtil.h delete mode 100644 third_party/mthreads/lib/Analysis/CMakeLists.txt delete mode 100644 third_party/mthreads/lib/CMakeLists.txt delete mode 100644 third_party/mthreads/lib/Conversion/CMakeLists.txt delete mode 100644 third_party/mthreads/lib/Conversion/TritonGPUToLLVM/AllocateSharedMemory.cpp delete mode 100644 third_party/mthreads/lib/Conversion/TritonGPUToLLVM/AllocateSharedMemoryUtility.cpp delete mode 100644 third_party/mthreads/lib/Conversion/TritonGPUToLLVM/CMakeLists.txt delete mode 100644 third_party/mthreads/lib/Conversion/TritonGPUToLLVM/ControlFlowOpToLLVM.cpp delete mode 100644 third_party/mthreads/lib/Conversion/TritonGPUToLLVM/DotOpToLLVM/FMADotUtility.cpp delete mode 100644 third_party/mthreads/lib/Conversion/TritonGPUToLLVM/MakeRangeOpToLLVM.cpp delete mode 100644 third_party/mthreads/lib/Conversion/TritonGPUToLLVM/PrintOpToLLVM.cpp delete mode 100644 third_party/mthreads/lib/Conversion/TritonGPUToLLVM/ReduceScanCommon.h delete mode 100644 third_party/mthreads/lib/Conversion/TritonGPUToLLVM/SPMDOpToLLVM.cpp delete mode 100644 third_party/mthreads/lib/Conversion/TritonGPUToLLVM/TypeConverter.cpp delete mode 100644 third_party/mthreads/lib/Conversion/TritonInstrumentToLLVM/CMakeLists.txt delete mode 100644 third_party/mthreads/lib/Conversion/TritonToTritonGPU/CMakeLists.txt delete mode 100644 third_party/mthreads/lib/Dialect/CMakeLists.txt delete mode 100644 third_party/mthreads/lib/Dialect/Gluon/CMakeLists.txt delete mode 100644 third_party/mthreads/lib/Dialect/Gluon/IR/CMakeLists.txt delete mode 100644 third_party/mthreads/lib/Dialect/Gluon/IR/Dialect.cpp delete mode 100644 third_party/mthreads/lib/Dialect/Gluon/Transforms/CMakeLists.txt delete mode 100644 third_party/mthreads/lib/Dialect/Gluon/Transforms/InferLayoutUtils.cpp delete mode 100644 third_party/mthreads/lib/Dialect/Gluon/Transforms/Inline.cpp delete mode 100644 third_party/mthreads/lib/Dialect/Gluon/Transforms/ResolveAutoEncodings.cpp delete mode 100644 third_party/mthreads/lib/Dialect/Gluon/Transforms/SimplifyControlFlow.cpp delete mode 100644 third_party/mthreads/lib/Dialect/NVGPU/CMakeLists.txt delete mode 100644 third_party/mthreads/lib/Dialect/NVWS/CMakeLists.txt delete mode 100644 third_party/mthreads/lib/Dialect/Triton/CMakeLists.txt delete mode 100644 third_party/mthreads/lib/Dialect/Triton/IR/CMakeLists.txt delete mode 100644 third_party/mthreads/lib/Dialect/Triton/IR/Canonicalize.td delete mode 100644 third_party/mthreads/lib/Dialect/Triton/IR/Dialect.cpp delete mode 100644 third_party/mthreads/lib/Dialect/Triton/IR/DiscardableAttributes.cpp delete mode 100644 third_party/mthreads/lib/Dialect/Triton/IR/OpInterfaces.cpp delete mode 100644 third_party/mthreads/lib/Dialect/Triton/Transforms/ArithTypeConversion.cpp delete mode 100644 third_party/mthreads/lib/Dialect/Triton/Transforms/CMakeLists.txt delete mode 100644 third_party/mthreads/lib/Dialect/Triton/Transforms/FunctionTypeConversion.cpp delete mode 100644 third_party/mthreads/lib/Dialect/Triton/Transforms/LoopAwareCSE.cpp delete mode 100644 third_party/mthreads/lib/Dialect/Triton/Transforms/LoopInvariantCodeMotion.cpp delete mode 100644 third_party/mthreads/lib/Dialect/Triton/Transforms/LoopPeeling.cpp delete mode 100644 third_party/mthreads/lib/Dialect/Triton/Transforms/ReorderBroadcast.cpp delete mode 100644 third_party/mthreads/lib/Dialect/TritonGPU/CMakeLists.txt delete mode 100644 third_party/mthreads/lib/Dialect/TritonGPU/IR/CMakeLists.txt delete mode 100644 third_party/mthreads/lib/Dialect/TritonGPU/Transforms/CMakeLists.txt delete mode 100644 third_party/mthreads/lib/Dialect/TritonGPU/Transforms/CombineTensorSelectAndIf.cpp delete mode 100644 third_party/mthreads/lib/Dialect/TritonGPU/Transforms/DecomposeScaledBlocked.cpp delete mode 100644 third_party/mthreads/lib/Dialect/TritonGPU/Transforms/FuseNestedLoops.cpp delete mode 100644 third_party/mthreads/lib/Dialect/TritonGPU/Transforms/HoistTMEMAlloc.cpp delete mode 100644 third_party/mthreads/lib/Dialect/TritonGPU/Transforms/LayoutPropagationUtility.cpp delete mode 100644 third_party/mthreads/lib/Dialect/TritonGPU/Transforms/Pipeliner/TestPipelineLowerLoop.cpp delete mode 100644 third_party/mthreads/lib/Dialect/TritonGPU/Transforms/ReorderInstructions.cpp delete mode 100644 third_party/mthreads/lib/Dialect/TritonGPU/Transforms/WarpSpecialization/PartitionBuilder.cpp delete mode 100644 third_party/mthreads/lib/Dialect/TritonInstrument/CMakeLists.txt delete mode 100644 third_party/mthreads/lib/Dialect/TritonInstrument/IR/CMakeLists.txt delete mode 100644 third_party/mthreads/lib/Dialect/TritonInstrument/IR/Dialect.cpp delete mode 100644 third_party/mthreads/lib/Dialect/TritonInstrument/IR/Ops.cpp delete mode 100644 third_party/mthreads/lib/Dialect/TritonInstrument/Transforms/CMakeLists.txt delete mode 100644 third_party/mthreads/lib/Dialect/TritonNvidiaGPU/CMakeLists.txt delete mode 100644 third_party/mthreads/lib/Dialect/TritonNvidiaGPU/IR/CMakeLists.txt delete mode 100644 third_party/mthreads/lib/Dialect/TritonNvidiaGPU/Transforms/CMakeLists.txt delete mode 100644 third_party/mthreads/lib/Dialect/TritonNvidiaGPU/Transforms/CheckMatmulTwoCTAs.cpp delete mode 100644 third_party/mthreads/lib/Target/CMakeLists.txt delete mode 100644 third_party/mthreads/lib/Target/LLVMIR/CMakeLists.txt delete mode 100644 third_party/mthreads/lib/Target/LLVMIR/LLVMIRBreakPhiStruct.cpp delete mode 100644 third_party/mthreads/lib/Target/LLVMIR/LLVMPasses.h delete mode 100644 third_party/mthreads/lib/Tools/CMakeLists.txt delete mode 100644 third_party/mthreads/lib/Tools/GenericSwizzling.cpp delete mode 100644 third_party/mthreads/lib/Tools/LayoutUtils.cpp delete mode 100644 third_party/mthreads/lib/Tools/LinearLayout.cpp rename third_party/mthreads/{ => spec_cpp}/include/triton/Analysis/Alias.h (100%) rename third_party/mthreads/{ => spec_cpp}/include/triton/Analysis/Allocation.h (100%) rename third_party/mthreads/{ => spec_cpp}/include/triton/Analysis/BufferRegion.h (100%) rename third_party/mthreads/{ => spec_cpp}/include/triton/Analysis/Membar.h (100%) rename third_party/mthreads/{ => spec_cpp}/include/triton/Analysis/Utility.h (100%) rename third_party/mthreads/{ => spec_cpp}/include/triton/Conversion/TritonGPUToLLVM/ElementwiseOpToLLVMBase.h (100%) rename third_party/mthreads/{ => spec_cpp}/include/triton/Conversion/TritonGPUToLLVM/PatternTritonGPUOpToLLVM.h (100%) rename third_party/mthreads/{ => spec_cpp}/include/triton/Conversion/TritonGPUToLLVM/TargetInfoBase.h (100%) rename third_party/mthreads/{ => spec_cpp}/include/triton/Conversion/TritonGPUToLLVM/Utility.h (100%) rename third_party/mthreads/{ => spec_cpp}/include/triton/Conversion/TritonGPUToLLVM/WarpSpecializeUtility.h (100%) rename third_party/mthreads/{ => spec_cpp}/include/triton/Dialect/Gluon/Transforms/Passes.td (100%) rename third_party/mthreads/{ => spec_cpp}/include/triton/Dialect/NVGPU/IR/CMakeLists.txt (100%) rename third_party/mthreads/{ => spec_cpp}/include/triton/Dialect/NVGPU/IR/Dialect.h (100%) rename third_party/mthreads/{ => spec_cpp}/include/triton/Dialect/NVGPU/IR/NVGPUAttrDefs.td (100%) rename third_party/mthreads/{ => spec_cpp}/include/triton/Dialect/NVGPU/IR/NVGPUDialect.td (100%) rename third_party/mthreads/{ => spec_cpp}/include/triton/Dialect/NVGPU/IR/NVGPUOps.td (100%) rename third_party/mthreads/{ => spec_cpp}/include/triton/Dialect/NVWS/IR/CMakeLists.txt (100%) rename third_party/mthreads/{ => spec_cpp}/include/triton/Dialect/NVWS/IR/Dialect.h (100%) rename third_party/mthreads/{ => spec_cpp}/include/triton/Dialect/NVWS/IR/NVWSAttrDefs.td (100%) rename third_party/mthreads/{ => spec_cpp}/include/triton/Dialect/NVWS/IR/NVWSDialect.td (100%) rename third_party/mthreads/{ => spec_cpp}/include/triton/Dialect/NVWS/IR/NVWSOpInterfaces.td (100%) rename third_party/mthreads/{ => spec_cpp}/include/triton/Dialect/NVWS/IR/NVWSOps.td (100%) rename third_party/mthreads/{ => spec_cpp}/include/triton/Dialect/NVWS/IR/NVWSTypes.td (100%) rename third_party/mthreads/{ => spec_cpp}/include/triton/Dialect/NVWS/Transforms/CMakeLists.txt (100%) rename third_party/mthreads/{ => spec_cpp}/include/triton/Dialect/NVWS/Transforms/Passes.h (100%) rename third_party/mthreads/{ => spec_cpp}/include/triton/Dialect/NVWS/Transforms/Passes.td (100%) create mode 100644 third_party/mthreads/spec_cpp/include/triton/Dialect/Triton/IR/CMakeLists.txt rename third_party/mthreads/{ => spec_cpp}/include/triton/Dialect/Triton/IR/Dialect.h (100%) rename third_party/mthreads/{ => spec_cpp}/include/triton/Dialect/Triton/IR/Traits.h (100%) rename third_party/mthreads/{ => spec_cpp}/include/triton/Dialect/Triton/IR/TritonInterfaces.td (100%) rename third_party/mthreads/{ => spec_cpp}/include/triton/Dialect/Triton/IR/TritonOps.td (100%) rename third_party/mthreads/{ => spec_cpp}/include/triton/Dialect/Triton/IR/TritonTypeInterfaces.td (100%) rename third_party/mthreads/{ => spec_cpp}/include/triton/Dialect/Triton/IR/TritonTypes.td (100%) rename third_party/mthreads/{ => spec_cpp}/include/triton/Dialect/Triton/IR/Types.h (100%) rename third_party/mthreads/{ => spec_cpp}/include/triton/Dialect/Triton/IR/Utility.h (100%) rename third_party/mthreads/{ => spec_cpp}/include/triton/Dialect/TritonGPU/IR/Attributes.h (100%) rename third_party/mthreads/{ => spec_cpp}/include/triton/Dialect/TritonGPU/IR/CGAEncodingAttr.h (100%) rename third_party/mthreads/{ => spec_cpp}/include/triton/Dialect/TritonGPU/IR/CGAEncodingAttr.td (100%) create mode 100644 third_party/mthreads/spec_cpp/include/triton/Dialect/TritonGPU/IR/CMakeLists.txt rename third_party/mthreads/{ => spec_cpp}/include/triton/Dialect/TritonGPU/IR/Dialect.h (100%) rename third_party/mthreads/{ => spec_cpp}/include/triton/Dialect/TritonGPU/IR/LinearLayoutConversions.h (100%) rename third_party/mthreads/{ => spec_cpp}/include/triton/Dialect/TritonGPU/IR/Traits.h (100%) rename third_party/mthreads/{ => spec_cpp}/include/triton/Dialect/TritonGPU/IR/TritonGPUAttrBase.td (100%) rename third_party/mthreads/{ => spec_cpp}/include/triton/Dialect/TritonGPU/IR/TritonGPUAttrDefs.td (100%) rename third_party/mthreads/{ => spec_cpp}/include/triton/Dialect/TritonGPU/IR/TritonGPUAttrImpls.td (100%) rename third_party/mthreads/{ => spec_cpp}/include/triton/Dialect/TritonGPU/IR/TritonGPUEnums.td (100%) rename third_party/mthreads/{ => spec_cpp}/include/triton/Dialect/TritonGPU/IR/TritonGPUInterfaces.h (100%) rename third_party/mthreads/{ => spec_cpp}/include/triton/Dialect/TritonGPU/IR/TritonGPUOps.td (100%) rename third_party/mthreads/{ => spec_cpp}/include/triton/Dialect/TritonGPU/IR/TritonGPUTypes.td (100%) rename third_party/mthreads/{ => spec_cpp}/include/triton/Dialect/TritonGPU/Transforms/CoalesceUtils.h (100%) rename third_party/mthreads/{ => spec_cpp}/include/triton/Dialect/TritonGPU/Transforms/PartitionSchedulingUtility.h (100%) rename third_party/mthreads/{ => spec_cpp}/include/triton/Dialect/TritonGPU/Transforms/Passes.h (100%) rename third_party/mthreads/{ => spec_cpp}/include/triton/Dialect/TritonGPU/Transforms/Passes.td (100%) rename third_party/mthreads/{ => spec_cpp}/include/triton/Dialect/TritonGPU/Transforms/Schedule.h (100%) rename third_party/mthreads/{ => spec_cpp}/include/triton/Dialect/TritonGPU/Transforms/TritonGPUConversion.h (100%) rename third_party/mthreads/{ => spec_cpp}/include/triton/Dialect/TritonGPU/Transforms/Utility.h (100%) rename third_party/mthreads/{ => spec_cpp}/include/triton/Dialect/TritonInstrument/IR/FunctionBuilder.h (100%) rename third_party/mthreads/{ => spec_cpp}/include/triton/Dialect/TritonInstrument/IR/TritonInstrumentOps.td (100%) rename third_party/mthreads/{ => spec_cpp}/include/triton/Dialect/TritonInstrument/IR/Utility.h (100%) create mode 100644 third_party/mthreads/spec_cpp/include/triton/Dialect/TritonNvidiaGPU/IR/CMakeLists.txt rename third_party/mthreads/{ => spec_cpp}/include/triton/Dialect/TritonNvidiaGPU/IR/Dialect.h (100%) rename third_party/mthreads/{ => spec_cpp}/include/triton/Dialect/TritonNvidiaGPU/IR/TritonNvidiaGPUAttrDefs.td (100%) rename third_party/mthreads/{ => spec_cpp}/include/triton/Dialect/TritonNvidiaGPU/IR/TritonNvidiaGPUDialect.td (100%) rename third_party/mthreads/{ => spec_cpp}/include/triton/Dialect/TritonNvidiaGPU/IR/TritonNvidiaGPUOpInterfaces.td (100%) rename third_party/mthreads/{ => spec_cpp}/include/triton/Dialect/TritonNvidiaGPU/IR/TritonNvidiaGPUOps.td (100%) rename third_party/mthreads/{ => spec_cpp}/include/triton/Dialect/TritonNvidiaGPU/IR/TritonNvidiaGPUTypes.td (100%) rename third_party/mthreads/{ => spec_cpp}/include/triton/Dialect/TritonNvidiaGPU/Transforms/TMAUtilities.h (100%) rename third_party/mthreads/{ => spec_cpp}/include/triton/Tools/LinearLayout.h (100%) rename third_party/mthreads/{ => spec_cpp}/include/triton/Tools/PluginUtils.h (100%) rename third_party/mthreads/{ => spec_cpp}/include/triton/Tools/Sys/GetEnv.hpp (100%) rename third_party/mthreads/{ => spec_cpp}/lib/Analysis/Alias.cpp (100%) rename third_party/mthreads/{ => spec_cpp}/lib/Analysis/Allocation.cpp (100%) rename third_party/mthreads/{ => spec_cpp}/lib/Analysis/AxisInfo.cpp (100%) rename third_party/mthreads/{ => spec_cpp}/lib/Analysis/BufferRegion.cpp (100%) rename third_party/mthreads/{ => spec_cpp}/lib/Analysis/Membar.cpp (100%) rename third_party/mthreads/{ => spec_cpp}/lib/Analysis/Utility.cpp (100%) rename third_party/mthreads/{ => spec_cpp}/lib/Conversion/TritonGPUToLLVM/AllocateWarpGroups.cpp (100%) rename third_party/mthreads/{ => spec_cpp}/lib/Conversion/TritonGPUToLLVM/AssertOpToLLVM.cpp (100%) rename third_party/mthreads/{ => spec_cpp}/lib/Conversion/TritonGPUToLLVM/ConvertLayoutOpToLLVM.cpp (100%) rename third_party/mthreads/{ => spec_cpp}/lib/Conversion/TritonGPUToLLVM/DotOpToLLVM/FMA.cpp (100%) rename third_party/mthreads/{ => spec_cpp}/lib/Conversion/TritonGPUToLLVM/ElementwiseOpToLLVM.cpp (100%) rename third_party/mthreads/{ => spec_cpp}/lib/Conversion/TritonGPUToLLVM/FuncOpToLLVM.cpp (100%) rename third_party/mthreads/{ => spec_cpp}/lib/Conversion/TritonGPUToLLVM/GatherOpToLLVM.cpp (100%) rename third_party/mthreads/{ => spec_cpp}/lib/Conversion/TritonGPUToLLVM/GlobalScratchMemoryAllocation.cpp (100%) rename third_party/mthreads/{ => spec_cpp}/lib/Conversion/TritonGPUToLLVM/HistogramOpToLLVM.cpp (100%) rename third_party/mthreads/{ => spec_cpp}/lib/Conversion/TritonGPUToLLVM/MemoryOpToLLVM.cpp (100%) rename third_party/mthreads/{ => spec_cpp}/lib/Conversion/TritonGPUToLLVM/ReduceOpToLLVM.cpp (100%) rename third_party/mthreads/{ => spec_cpp}/lib/Conversion/TritonGPUToLLVM/ScanOpToLLVM.cpp (100%) rename third_party/mthreads/{ => spec_cpp}/lib/Conversion/TritonGPUToLLVM/Utility.cpp (100%) rename third_party/mthreads/{ => spec_cpp}/lib/Conversion/TritonGPUToLLVM/ViewOpToLLVM.cpp (100%) rename third_party/mthreads/{ => spec_cpp}/lib/Conversion/TritonGPUToLLVM/WarpSpecializeUtility.cpp (100%) rename third_party/mthreads/{ => spec_cpp}/lib/Conversion/TritonInstrumentToLLVM/InstrumentationToLLVM.cpp (100%) rename third_party/mthreads/{ => spec_cpp}/lib/Conversion/TritonToTritonGPU/RelayoutTritonGPU.cpp (100%) rename third_party/mthreads/{ => spec_cpp}/lib/Conversion/TritonToTritonGPU/TritonGPUConversion.cpp (100%) rename third_party/mthreads/{ => spec_cpp}/lib/Conversion/TritonToTritonGPU/TritonToTritonGPUPass.cpp (100%) rename third_party/mthreads/{ => spec_cpp}/lib/Dialect/Gluon/Transforms/Canonicalize.cpp (100%) rename third_party/mthreads/{ => spec_cpp}/lib/Dialect/Gluon/Transforms/InferCoalescedEncodings.cpp (100%) rename third_party/mthreads/{ => spec_cpp}/lib/Dialect/NVGPU/IR/CMakeLists.txt (100%) rename third_party/mthreads/{ => spec_cpp}/lib/Dialect/NVGPU/IR/Dialect.cpp (100%) rename third_party/mthreads/{ => spec_cpp}/lib/Dialect/NVWS/IR/CMakeLists.txt (100%) rename third_party/mthreads/{ => spec_cpp}/lib/Dialect/NVWS/IR/Dialect.cpp (100%) rename third_party/mthreads/{ => spec_cpp}/lib/Dialect/NVWS/IR/Ops.cpp (100%) rename third_party/mthreads/{ => spec_cpp}/lib/Dialect/NVWS/Transforms/AssignStagePhase.cpp (100%) rename third_party/mthreads/{ => spec_cpp}/lib/Dialect/NVWS/Transforms/CMakeLists.txt (100%) rename third_party/mthreads/{ => spec_cpp}/lib/Dialect/NVWS/Transforms/HoistTmemStore.cpp (100%) rename third_party/mthreads/{ => spec_cpp}/lib/Dialect/NVWS/Transforms/InsertAref.cpp (100%) rename third_party/mthreads/{ => spec_cpp}/lib/Dialect/NVWS/Transforms/InsertTmemAref.cpp (100%) rename third_party/mthreads/{ => spec_cpp}/lib/Dialect/NVWS/Transforms/LowerAref.cpp (100%) rename third_party/mthreads/{ => spec_cpp}/lib/Dialect/NVWS/Transforms/LowerWarpGroup.cpp (100%) rename third_party/mthreads/{ => spec_cpp}/lib/Dialect/NVWS/Transforms/Utilities.cpp (100%) rename third_party/mthreads/{ => spec_cpp}/lib/Dialect/NVWS/Transforms/Utilities.h (100%) rename third_party/mthreads/{ => spec_cpp}/lib/Dialect/Triton/IR/Ops.cpp (100%) rename third_party/mthreads/{ => spec_cpp}/lib/Dialect/Triton/IR/Traits.cpp (100%) rename third_party/mthreads/{ => spec_cpp}/lib/Dialect/Triton/IR/Types.cpp (100%) rename third_party/mthreads/{ => spec_cpp}/lib/Dialect/Triton/IR/Utility.cpp (100%) rename third_party/mthreads/{ => spec_cpp}/lib/Dialect/Triton/Transforms/Combine.cpp (100%) rename third_party/mthreads/{ => spec_cpp}/lib/Dialect/Triton/Transforms/Combine.td (100%) rename third_party/mthreads/{ => spec_cpp}/lib/Dialect/Triton/Transforms/LoopUnroll.cpp (100%) rename third_party/mthreads/{ => spec_cpp}/lib/Dialect/Triton/Transforms/RewriteTensorDescriptorToPointer.cpp (100%) rename third_party/mthreads/{ => spec_cpp}/lib/Dialect/Triton/Transforms/RewriteTensorPointer.cpp (100%) rename third_party/mthreads/{ => spec_cpp}/lib/Dialect/TritonGPU/IR/Dialect.cpp (100%) rename third_party/mthreads/{ => spec_cpp}/lib/Dialect/TritonGPU/IR/LinearLayoutConversions.cpp (100%) rename third_party/mthreads/{ => spec_cpp}/lib/Dialect/TritonGPU/IR/Ops.cpp (100%) rename third_party/mthreads/{ => spec_cpp}/lib/Dialect/TritonGPU/IR/Types.cpp (100%) rename third_party/mthreads/{ => spec_cpp}/lib/Dialect/TritonGPU/Transforms/AccelerateMatmul.cpp (100%) rename third_party/mthreads/{ => spec_cpp}/lib/Dialect/TritonGPU/Transforms/Coalesce.cpp (100%) rename third_party/mthreads/{ => spec_cpp}/lib/Dialect/TritonGPU/Transforms/CoalesceAsyncCopy.cpp (100%) rename third_party/mthreads/{ => spec_cpp}/lib/Dialect/TritonGPU/Transforms/CoalesceUtils.cpp (100%) rename third_party/mthreads/{ => spec_cpp}/lib/Dialect/TritonGPU/Transforms/F32DotTC.cpp (100%) rename third_party/mthreads/{ => spec_cpp}/lib/Dialect/TritonGPU/Transforms/OptimizeAccumulatorInit.cpp (100%) rename third_party/mthreads/{ => spec_cpp}/lib/Dialect/TritonGPU/Transforms/OptimizeDotOperands.cpp (100%) rename third_party/mthreads/{ => spec_cpp}/lib/Dialect/TritonGPU/Transforms/OptimizeThreadLocality.cpp (100%) rename third_party/mthreads/{ => spec_cpp}/lib/Dialect/TritonGPU/Transforms/Pipeliner/AssignLatencies.cpp (100%) rename third_party/mthreads/{ => spec_cpp}/lib/Dialect/TritonGPU/Transforms/Pipeliner/LowerLoops.cpp (100%) rename third_party/mthreads/{ => spec_cpp}/lib/Dialect/TritonGPU/Transforms/Pipeliner/MMAv5PipelineUtility.cpp (100%) rename third_party/mthreads/{ => spec_cpp}/lib/Dialect/TritonGPU/Transforms/Pipeliner/PipelineExpander.cpp (100%) rename third_party/mthreads/{ => spec_cpp}/lib/Dialect/TritonGPU/Transforms/Pipeliner/PipeliningUtility.cpp (100%) rename third_party/mthreads/{ => spec_cpp}/lib/Dialect/TritonGPU/Transforms/Pipeliner/Schedule.cpp (100%) rename third_party/mthreads/{ => spec_cpp}/lib/Dialect/TritonGPU/Transforms/Pipeliner/ScheduleLoops.cpp (100%) rename third_party/mthreads/{ => spec_cpp}/lib/Dialect/TritonGPU/Transforms/Pipeliner/SoftwarePipeliner.cpp (100%) rename third_party/mthreads/{ => spec_cpp}/lib/Dialect/TritonGPU/Transforms/Pipeliner/TMAStoresPipeline.cpp (100%) rename third_party/mthreads/{ => spec_cpp}/lib/Dialect/TritonGPU/Transforms/Pipeliner/WGMMAPipeline.cpp (100%) rename third_party/mthreads/{ => spec_cpp}/lib/Dialect/TritonGPU/Transforms/Prefetch.cpp (100%) rename third_party/mthreads/{ => spec_cpp}/lib/Dialect/TritonGPU/Transforms/ReduceDataDuplication.cpp (100%) rename third_party/mthreads/{ => spec_cpp}/lib/Dialect/TritonGPU/Transforms/RemoveLayoutConversions.cpp (100%) rename third_party/mthreads/{ => spec_cpp}/lib/Dialect/TritonGPU/Transforms/Utility.cpp (100%) rename third_party/mthreads/{ => spec_cpp}/lib/Dialect/TritonGPU/Transforms/WarpSpecialization/AutomaticWarpSpecialization.cpp (100%) rename third_party/mthreads/{ => spec_cpp}/lib/Dialect/TritonGPU/Transforms/WarpSpecialization/OptimizePartitionWarps.cpp (100%) rename third_party/mthreads/{ => spec_cpp}/lib/Dialect/TritonGPU/Transforms/WarpSpecialization/Partition.cpp (100%) rename third_party/mthreads/{ => spec_cpp}/lib/Dialect/TritonGPU/Transforms/WarpSpecialization/PartitionLoops.cpp (100%) rename third_party/mthreads/{ => spec_cpp}/lib/Dialect/TritonGPU/Transforms/WarpSpecialization/PartitionScheduling.cpp (100%) rename third_party/mthreads/{ => spec_cpp}/lib/Dialect/TritonGPU/Transforms/WarpSpecialization/PartitionSchedulingUtility.cpp (100%) rename third_party/mthreads/{ => spec_cpp}/lib/Dialect/TritonInstrument/IR/FunctionBuilder.cpp (100%) rename third_party/mthreads/{ => spec_cpp}/lib/Dialect/TritonInstrument/IR/Utility.cpp (100%) rename third_party/mthreads/{ => spec_cpp}/lib/Dialect/TritonInstrument/Transforms/ConcurrencySanitizer.cpp (100%) rename third_party/mthreads/{ => spec_cpp}/lib/Dialect/TritonNvidiaGPU/IR/Dialect.cpp (100%) rename third_party/mthreads/{ => spec_cpp}/lib/Dialect/TritonNvidiaGPU/IR/Ops.cpp (100%) rename third_party/mthreads/{ => spec_cpp}/lib/Dialect/TritonNvidiaGPU/IR/TensorMemoryUtils.cpp (100%) rename third_party/mthreads/{ => spec_cpp}/lib/Dialect/TritonNvidiaGPU/Transforms/FenceInsertion.cpp (100%) rename third_party/mthreads/{ => spec_cpp}/lib/Dialect/TritonNvidiaGPU/Transforms/InterleaveTMem.cpp (100%) rename third_party/mthreads/{ => spec_cpp}/lib/Dialect/TritonNvidiaGPU/Transforms/MMALowering.cpp (100%) rename third_party/mthreads/{ => spec_cpp}/lib/Dialect/TritonNvidiaGPU/Transforms/OptimizeDescriptorEncoding.cpp (100%) rename third_party/mthreads/{ => spec_cpp}/lib/Dialect/TritonNvidiaGPU/Transforms/OptimizeTMemLayouts.cpp (100%) rename third_party/mthreads/{ => spec_cpp}/lib/Dialect/TritonNvidiaGPU/Transforms/PlanCTA.cpp (100%) rename third_party/mthreads/{ => spec_cpp}/lib/Dialect/TritonNvidiaGPU/Transforms/PromoteLHSToTMem.cpp (100%) rename third_party/mthreads/{ => spec_cpp}/lib/Dialect/TritonNvidiaGPU/Transforms/ProxyFenceInsertion.cpp (100%) rename third_party/mthreads/{ => spec_cpp}/lib/Dialect/TritonNvidiaGPU/Transforms/RemoveTMEMTokens.cpp (100%) rename third_party/mthreads/{ => spec_cpp}/lib/Dialect/TritonNvidiaGPU/Transforms/TMALowering.cpp (100%) rename third_party/mthreads/{ => spec_cpp}/lib/Dialect/TritonNvidiaGPU/Transforms/TMAUtilities.cpp (100%) rename third_party/mthreads/{ => spec_cpp}/lib/Dialect/TritonNvidiaGPU/Transforms/TensorMemoryAllocation.cpp (100%) rename third_party/mthreads/{ => spec_cpp}/lib/Target/LLVMIR/LLVMDILocalVariable.cpp (100%) rename third_party/mthreads/{ => spec_cpp}/lib/Target/LLVMIR/LLVMDIScope.cpp (100%) rename third_party/mthreads/{ => spec_cpp}/lib/Target/LLVMIR/LLVMDIUtils.cpp (100%) rename third_party/mthreads/{ => spec_cpp}/lib/Target/LLVMIR/LLVMDIUtils.h (100%) rename third_party/mthreads/{ => spec_cpp}/lib/Tools/PluginUtils.cpp (100%) diff --git a/CMakeLists.txt b/CMakeLists.txt index f2c9b04193..5cec1c4b92 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -375,6 +375,11 @@ if(NOT TRITON_BUILD_PYTHON_MODULE) add_subdirectory(third_party/tle) flagtree_add_tle_generated_header_dependencies() endif() + # flagtree backend cmake specialization + if(FLAGTREE_BACKEND) + include(FlagTreeBackendSpec) + flagtree_apply_backend_source_overrides("${FLAGTREE_BACKEND_DIR}") + endif() endif() find_package(Threads REQUIRED) diff --git a/cmake/FlagTreeBackendSpec.cmake b/cmake/FlagTreeBackendSpec.cmake index bb0bb40b12..4f48f979c6 100644 --- a/cmake/FlagTreeBackendSpec.cmake +++ b/cmake/FlagTreeBackendSpec.cmake @@ -145,9 +145,26 @@ function(flagtree_apply_backend_source_overrides backend_root) list(LENGTH _spec_sources_in_core_root _candidate_count) if(_candidate_count EQUAL 0) - message(FATAL_ERROR - "Backend spec source ${_spec_source} has no owner target in the " - "configured core roots for mirrored main source ${_relative_path}") + set(_added_owner_targets) + foreach(_index RANGE 0 ${_last_source_index}) + list(GET _source_index ${_index} _indexed_source) + if(_indexed_source STREQUAL "${_spec_source}") + list(GET _target_index ${_index} _added_owner_target) + list(APPEND _added_owner_targets "${_added_owner_target}") + endif() + endforeach() + list(REMOVE_DUPLICATES _added_owner_targets) + list(LENGTH _added_owner_targets _added_owner_count) + if(_added_owner_count EQUAL 0) + message(FATAL_ERROR + "Backend spec source ${_spec_source} has no mirrored main source " + "or explicit owner target for ${_relative_path}") + endif() + foreach(_added_owner_target IN LISTS _added_owner_targets) + message(STATUS + "SPEC ADD: ${_relative_path} -> ${_added_owner_target}") + endforeach() + continue() elseif(_candidate_count GREATER 1) message(FATAL_ERROR "Backend spec source ${_spec_source} matches multiple preferred main " diff --git a/cmake/FlagTreeOptions.cmake b/cmake/FlagTreeOptions.cmake index 531ae89af7..852d522ce2 100644 --- a/cmake/FlagTreeOptions.cmake +++ b/cmake/FlagTreeOptions.cmake @@ -109,7 +109,9 @@ macro(flagtree_include_directories root_include_dir) # flagtree spec include dir set(BACKEND_SPEC_INCLUDE_DIR ${FLAGTREE_BACKEND_DIR}/spec_cpp/include) if(FLAGTREE_BACKEND AND EXISTS ${BACKEND_SPEC_INCLUDE_DIR}) - include_directories(${BACKEND_SPEC_INCLUDE_DIR}) + include_directories(BEFORE + ${BACKEND_SPEC_INCLUDE_DIR} + ${PROJECT_BINARY_DIR}/third_party/${FLAGTREE_BACKEND}/spec_cpp/include) endif() # flagtree third_party include dir @@ -158,9 +160,26 @@ macro(flagtree_configure_core_source) endif() if(FLAGTREE_BACKEND MATCHES - "^(xpu|cambricon|aipu|tsingmicro|enflame|rpu|thrive|tileir|ppu)$") + "^(xpu|cambricon|aipu|tsingmicro|enflame|rpu|thrive|mthreads|tileir|ppu)$") include_directories(${PROJECT_SOURCE_DIR}/include) include_directories(${PROJECT_BINARY_DIR}/include) # Tablegen'd files + if(FLAGTREE_BACKEND STREQUAL "mthreads") + set(_flagtree_mthreads_root + ${PROJECT_SOURCE_DIR}/third_party/mthreads) + include_directories(BEFORE + ${_flagtree_mthreads_root}/spec_cpp/include + ${PROJECT_BINARY_DIR}/third_party/mthreads/spec_cpp/include + ${_flagtree_mthreads_root}/spec_cpp + ${PROJECT_BINARY_DIR}/third_party/mthreads/spec_cpp + ${_flagtree_mthreads_root} + ${PROJECT_BINARY_DIR}/third_party/mthreads + ${_flagtree_mthreads_root}/musa/include + ${PROJECT_BINARY_DIR}/third_party/mthreads/musa/include + ${_flagtree_mthreads_root}/tle/dialect/include + ${PROJECT_BINARY_DIR}/third_party/mthreads/tle/dialect/include + ${_flagtree_mthreads_root}/tle/frontend/include + ${PROJECT_BINARY_DIR}/third_party/mthreads/tle/frontend/include) + endif() if(FLAGTREE_BACKEND STREQUAL "xpu") include_directories(${PROJECT_SOURCE_DIR}/third_party/nvidia/include) include_directories(${PROJECT_BINARY_DIR}/third_party/nvidia/include) # Tablegen'd files diff --git a/include/triton/Dialect/Gluon/Transforms/CMakeLists.txt b/include/triton/Dialect/Gluon/Transforms/CMakeLists.txt index a2d298d0c1..44aec47fd8 100644 --- a/include/triton/Dialect/Gluon/Transforms/CMakeLists.txt +++ b/include/triton/Dialect/Gluon/Transforms/CMakeLists.txt @@ -1,3 +1,3 @@ -set(LLVM_TARGET_DEFINITIONS Passes.td) +flagtree_spec_td_set(LLVM_TARGET_DEFINITIONS Passes.td) mlir_tablegen(Passes.h.inc -gen-pass-decls -name Gluon) add_public_tablegen_target(GluonTransformsIncGen) diff --git a/include/triton/Dialect/Triton/IR/CMakeLists.txt b/include/triton/Dialect/Triton/IR/CMakeLists.txt index bfe2237609..daece64408 100644 --- a/include/triton/Dialect/Triton/IR/CMakeLists.txt +++ b/include/triton/Dialect/Triton/IR/CMakeLists.txt @@ -12,11 +12,11 @@ mlir_tablegen(Dialect.h.inc -gen-dialect-decls) mlir_tablegen(Dialect.cpp.inc -gen-dialect-defs) add_mlir_doc(TritonDialect TritonDialect dialects/ -gen-dialect-doc) -set(LLVM_TARGET_DEFINITIONS TritonTypes.td) +flagtree_spec_td_set(LLVM_TARGET_DEFINITIONS TritonTypes.td) mlir_tablegen(Types.h.inc -gen-typedef-decls) mlir_tablegen(Types.cpp.inc -gen-typedef-defs) -set(LLVM_TARGET_DEFINITIONS TritonInterfaces.td) +flagtree_spec_td_set(LLVM_TARGET_DEFINITIONS TritonInterfaces.td) mlir_tablegen(AttrInterfaces.h.inc -gen-attr-interface-decls) mlir_tablegen(AttrInterfaces.cpp.inc -gen-attr-interface-defs) diff --git a/include/triton/Dialect/TritonGPU/IR/CMakeLists.txt b/include/triton/Dialect/TritonGPU/IR/CMakeLists.txt index cafd9a38ff..c1d5946467 100644 --- a/include/triton/Dialect/TritonGPU/IR/CMakeLists.txt +++ b/include/triton/Dialect/TritonGPU/IR/CMakeLists.txt @@ -1,6 +1,6 @@ set(MLIR_BINARY_DIR ${CMAKE_BINARY_DIR}) -set(LLVM_TARGET_DEFINITIONS TritonGPUOps.td) +flagtree_spec_td_set(LLVM_TARGET_DEFINITIONS TritonGPUOps.td) mlir_tablegen(Dialect.h.inc -gen-dialect-decls -dialect=ttg) mlir_tablegen(Dialect.cpp.inc -gen-dialect-defs -dialect=ttg) mlir_tablegen(Ops.h.inc -gen-op-decls) @@ -18,7 +18,7 @@ mlir_tablegen(AttrDefs.h.inc -gen-attrdef-decls) mlir_tablegen(OpsEnums.h.inc -gen-enum-decls) mlir_tablegen(OpsEnums.cpp.inc -gen-enum-defs) -set(LLVM_TARGET_DEFINITIONS TritonGPUAttrImpls.td) +flagtree_spec_td_set(LLVM_TARGET_DEFINITIONS TritonGPUAttrImpls.td) mlir_tablegen(AttrDefs.cpp.inc -gen-attrdef-defs) add_public_tablegen_target(TritonGPUAttrDefsIncGen) diff --git a/include/triton/Dialect/TritonInstrument/IR/CMakeLists.txt b/include/triton/Dialect/TritonInstrument/IR/CMakeLists.txt index 2af09f9046..fcf1486b00 100644 --- a/include/triton/Dialect/TritonInstrument/IR/CMakeLists.txt +++ b/include/triton/Dialect/TritonInstrument/IR/CMakeLists.txt @@ -5,7 +5,7 @@ mlir_tablegen(Dialect.h.inc -gen-dialect-decls -dialect=tti) mlir_tablegen(Dialect.cpp.inc -gen-dialect-defs -dialect=tti) add_mlir_doc(TritonInstrumentDialect TritonInstrumentDialect dialects/ -gen-dialect-doc) -set(LLVM_TARGET_DEFINITIONS TritonInstrumentOps.td) +flagtree_spec_td_set(LLVM_TARGET_DEFINITIONS TritonInstrumentOps.td) mlir_tablegen(Ops.h.inc -gen-op-decls) mlir_tablegen(Ops.cpp.inc -gen-op-defs) mlir_tablegen(OpsEnums.h.inc -gen-enum-decls) diff --git a/include/triton/Dialect/TritonNvidiaGPU/IR/CMakeLists.txt b/include/triton/Dialect/TritonNvidiaGPU/IR/CMakeLists.txt index 7cb25f5044..e36d751511 100644 --- a/include/triton/Dialect/TritonNvidiaGPU/IR/CMakeLists.txt +++ b/include/triton/Dialect/TritonNvidiaGPU/IR/CMakeLists.txt @@ -1,6 +1,6 @@ set(MLIR_BINARY_DIR ${CMAKE_BINARY_DIR}) -set(LLVM_TARGET_DEFINITIONS TritonNvidiaGPUOps.td) +flagtree_spec_td_set(LLVM_TARGET_DEFINITIONS TritonNvidiaGPUOps.td) mlir_tablegen(Dialect.h.inc -gen-dialect-decls -dialect=ttng) mlir_tablegen(Dialect.cpp.inc -gen-dialect-defs -dialect=ttng) mlir_tablegen(Ops.h.inc -gen-op-decls) @@ -9,14 +9,14 @@ add_mlir_doc(TritonNvidiaGPUDialect TritonNvidiaGPUDialect dialects/ -gen-dialec add_mlir_doc(TritonNvidiaGPUOps TritonNvidiaGPUOps dialects/ -gen-op-doc) add_public_tablegen_target(TritonNvidiaGPUTableGen) -set(LLVM_TARGET_DEFINITIONS TritonNvidiaGPUAttrDefs.td) +flagtree_spec_td_set(LLVM_TARGET_DEFINITIONS TritonNvidiaGPUAttrDefs.td) mlir_tablegen(TritonNvidiaGPUAttrDefs.h.inc -gen-attrdef-decls) mlir_tablegen(TritonNvidiaGPUAttrDefs.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(TritonNvidiaGPUAttrDefsIncGen) -set(LLVM_TARGET_DEFINITIONS TritonNvidiaGPUOpInterfaces.td) +flagtree_spec_td_set(LLVM_TARGET_DEFINITIONS TritonNvidiaGPUOpInterfaces.td) mlir_tablegen(TritonNvidiaGPUOpInterfaces.h.inc -gen-op-interface-decls) mlir_tablegen(TritonNvidiaGPUOpInterfaces.cpp.inc -gen-op-interface-defs) add_public_tablegen_target(TritonNvidiaGPUOpInterfacesIncGen) diff --git a/lib/Dialect/Triton/Transforms/CMakeLists.txt b/lib/Dialect/Triton/Transforms/CMakeLists.txt index ef7187b8f5..7b82c59a2b 100644 --- a/lib/Dialect/Triton/Transforms/CMakeLists.txt +++ b/lib/Dialect/Triton/Transforms/CMakeLists.txt @@ -1,4 +1,4 @@ -set(LLVM_TARGET_DEFINITIONS Combine.td) +flagtree_spec_td_set(LLVM_TARGET_DEFINITIONS Combine.td) mlir_tablegen(TritonCombine.inc -gen-rewriters) add_public_tablegen_target(TritonCombineIncGen) diff --git a/third_party/mthreads/CMakeLists.txt b/third_party/mthreads/CMakeLists.txt index d7bcef927e..43da609274 100644 --- a/third_party/mthreads/CMakeLists.txt +++ b/third_party/mthreads/CMakeLists.txt @@ -1,7 +1,9 @@ include_directories(${CMAKE_CURRENT_SOURCE_DIR}) include_directories(${CMAKE_CURRENT_BINARY_DIR}) -include_directories(${CMAKE_CURRENT_SOURCE_DIR}/include) -include_directories(${CMAKE_CURRENT_BINARY_DIR}/include) +include_directories(${CMAKE_CURRENT_SOURCE_DIR}/spec_cpp) +include_directories(${CMAKE_CURRENT_BINARY_DIR}/spec_cpp) +include_directories(${CMAKE_CURRENT_SOURCE_DIR}/spec_cpp/include) +include_directories(${CMAKE_CURRENT_BINARY_DIR}/spec_cpp/include) include_directories(${CMAKE_CURRENT_SOURCE_DIR}/musa/include) include_directories(${CMAKE_CURRENT_BINARY_DIR}/musa/include) include_directories(${CMAKE_CURRENT_SOURCE_DIR}/tle/dialect/include) @@ -19,9 +21,58 @@ if(FLAGTREE_MTHREADS_TLE) ${CMAKE_SOURCE_DIR}/third_party/tle/dialect/lib/IR ${CMAKE_BINARY_DIR}/third_party/tle/dialect/lib/IR) add_subdirectory(tle) + add_dependencies(TritonToTritonGPU MUSATLETableGen) + target_link_libraries(TritonToTritonGPU PUBLIC MUSATLEIR) endif() -add_subdirectory(include) -add_subdirectory(lib) + +add_subdirectory(spec_cpp/include/triton/Dialect/Triton/IR) +add_subdirectory(spec_cpp/include/triton/Dialect/TritonGPU/IR) +add_subdirectory(spec_cpp/include/triton/Dialect/TritonNvidiaGPU/IR) +add_subdirectory(spec_cpp/include/triton/Dialect/NVGPU/IR) +add_subdirectory(spec_cpp/include/triton/Dialect/NVWS/IR) +add_subdirectory(spec_cpp/include/triton/Dialect/NVWS/Transforms) + +add_subdirectory(spec_cpp/lib/Dialect/NVGPU/IR) +add_subdirectory(spec_cpp/lib/Dialect/NVWS/IR) +add_subdirectory(spec_cpp/lib/Dialect/NVWS/Transforms) + +set(MTHREADS_SPEC_CPP_LIB + ${CMAKE_CURRENT_SOURCE_DIR}/spec_cpp/lib) + +target_sources(TritonAnalysis PRIVATE + ${MTHREADS_SPEC_CPP_LIB}/Analysis/BufferRegion.cpp) +target_sources(TritonGPUToLLVM PRIVATE + ${MTHREADS_SPEC_CPP_LIB}/Conversion/TritonGPUToLLVM/WarpSpecializeUtility.cpp) +target_sources(TritonGPUTransforms PRIVATE + ${MTHREADS_SPEC_CPP_LIB}/Dialect/TritonGPU/Transforms/WarpSpecialization/PartitionSchedulingUtility.cpp) +target_sources(TritonNvidiaGPUTransforms PRIVATE + ${MTHREADS_SPEC_CPP_LIB}/Dialect/TritonNvidiaGPU/Transforms/ProxyFenceInsertion.cpp) +target_sources(TritonLLVMIR PRIVATE + ${MTHREADS_SPEC_CPP_LIB}/Target/LLVMIR/LLVMDIUtils.cpp) +target_sources(TritonTools PRIVATE + ${MTHREADS_SPEC_CPP_LIB}/Tools/PluginUtils.cpp) + +set_source_files_properties( + ${PROJECT_SOURCE_DIR}/lib/Dialect/TritonGPU/Transforms/ProcessSharedMemoryHint.cpp + TARGET_DIRECTORY TritonGPUTransforms + PROPERTIES HEADER_FILE_ONLY ON) +set_source_files_properties( + ${PROJECT_SOURCE_DIR}/lib/Dialect/TritonNvidiaGPU/Transforms/ProxFenceInsertion.cpp + TARGET_DIRECTORY TritonNvidiaGPUTransforms + PROPERTIES HEADER_FILE_ONLY ON) + +add_dependencies(TritonTableGen TritonTypeInterfacesIncGen) +add_dependencies(TritonGPUTableGen TritonGPUOpsEnumsIncGen) +add_dependencies(TritonGPUAttrDefsIncGen + TritonGPUCGAAttrIncGen + TritonGPUOpsEnumsIncGen) +add_dependencies(TritonNvidiaGPUTableGen TritonNvidiaGPUTypesIncGen) + +if(TARGET PrintLoadStoreMemSpaces) + set_target_properties(PrintLoadStoreMemSpaces + PROPERTIES EXCLUDE_FROM_ALL TRUE) +endif() + add_subdirectory(musa) if(TRITON_BUILD_PYTHON_MODULE) if(FLAGTREE_MTHREADS_TLE) diff --git a/third_party/mthreads/include/CMakeLists.txt b/third_party/mthreads/include/CMakeLists.txt deleted file mode 100644 index 109c292fea..0000000000 --- a/third_party/mthreads/include/CMakeLists.txt +++ /dev/null @@ -1 +0,0 @@ -add_subdirectory(triton) diff --git a/third_party/mthreads/include/triton/Analysis/AxisInfo.h b/third_party/mthreads/include/triton/Analysis/AxisInfo.h deleted file mode 100644 index f252081f43..0000000000 --- a/third_party/mthreads/include/triton/Analysis/AxisInfo.h +++ /dev/null @@ -1,271 +0,0 @@ -#ifndef TRITON_ANALYSIS_AXISINFO_H -#define TRITON_ANALYSIS_AXISINFO_H - -#include "mlir/Analysis/DataFlow/SparseAnalysis.h" -#include "llvm/Support/raw_ostream.h" - -#include "mlir/Support/LLVM.h" -#include "triton/Analysis/Utility.h" - -#include - -namespace mlir::triton { - -//===----------------------------------------------------------------------===// -// AxisInfo -//===----------------------------------------------------------------------===// - -/// This lattice value represents known information on the axes of a lattice. -class AxisInfo { -public: - typedef SmallVector DimVectorT; - -public: - AxisInfo() : AxisInfo({}, {}, {}) {} - - AxisInfo(ArrayRef contiguity, ArrayRef divisibility, - ArrayRef constancy) - : AxisInfo(contiguity, divisibility, constancy, std::nullopt) {} - - AxisInfo(ArrayRef contiguity, ArrayRef divisibility, - ArrayRef constancy, std::optional constantValue) - : contiguity(contiguity), divisibility(divisibility), - constancy(constancy), constantValue(constantValue) { - assert(divisibility.size() == contiguity.size()); - assert(constancy.size() == contiguity.size()); - } - - // contiguity[d] is the length of the shortest sequence of contiguous integers - // along dimension d. - // - // If we have an array of N elements with a contiguity value C, then the array - // can be divided into a list of N/C sequences of C contiguous elements. - // Since we have N = 2^k, C must be a power of two. - // - // For example, the 2D array - // - // [[10, 11, 12, 13, 18, 19, 20, 21], - // [20, 21, 22, 23, 28, 29, 30, 31]] - // - // has contiguity [1, 4], and - // - // [[12, 16, 20, 24], - // [13, 17, 21, 25], - // [14, 18, 22, 26], - // [15, 19, 23, 27], - // [18, 22, 26, 30], - // [19, 23, 27, 31]] - // - // has contiguity [2, 1]. - int64_t getContiguity(size_t dim) const { return contiguity[dim]; } - const DimVectorT &getContiguity() const { return contiguity; } - - // divisibility[d] is the largest power of two that divides the first element - // of all groups of length contiguity[d] along dimension d. - // - // For example, - // - // [[10, 11, 12, 13, 18, 19, 20, 21], - // [20, 21, 22, 23, 28, 29, 30, 31]] - // - // has divisibility [1, 2], and - // - // [[12, 16, 20, 24], - // [13, 17, 21, 25], - // [14, 18, 22, 26], - // [15, 19, 23, 27]] - // - // has divisibility [4, 1]. - // - // On the other hand, - // - // [0, 1, 2, 0, 4, 5, 6, 7] - // - // has divisibility 1 because its contiguity is 1. - int64_t getDivisibility(size_t dim) const { return divisibility[dim]; } - const DimVectorT &getDivisibility() const { return divisibility; } - - // constancy[d] is the length of the shortest sequence of repeating integers - // along dimension d. - // - // This is particularly useful to infer the contiguity of operations (e.g. - // add) involving a constant. - // - // If we have an array of N elements, with a constancy value C, then the array - // can be divided into a list of N/C sequences of C elements with the same - // value. Since we have N = 2^k, C must be a power of two. - // - // For example - // - // [[8, 8, 8, 8, 12, 12, 12, 12], - // [16, 16, 16, 16, 20, 20, 20, 20]] - // - // has constancy [1, 4]. - int64_t getConstancy(size_t dim) const { return constancy[dim]; } - const DimVectorT &getConstancy() const { return constancy; } - - int getRank() const { return contiguity.size(); } - - std::optional getConstantValue() const { return constantValue; } - - static void initPessimisticStateFromFunc(int argNumber, - FunctionOpInterface funcOp, - DimVectorT *contiguity, - DimVectorT *divisibility, - DimVectorT *constancy); - - static void initDimVectorFromHint(Attribute attr, DimVectorT *vec); - - bool operator==(const AxisInfo &other) const { - return contiguity == other.contiguity && - divisibility == other.divisibility && constancy == other.constancy && - constantValue == other.constantValue; - } - - static AxisInfo getPessimisticValueState(Value value); - - // The gcd of both arguments for each dimension - static AxisInfo join(const AxisInfo &lhs, const AxisInfo &rhs); - - void print(raw_ostream &os) const { - auto print = [&](StringRef name, DimVectorT vec) { - os << name << " = ["; - llvm::interleaveComma(vec, os); - os << "]"; - }; - print("contiguity", contiguity); - print(", divisibility", divisibility); - print(", constancy", constancy); - os << ", constant_value = "; - if (constantValue) - os << *constantValue; - else - os << ""; - } - -private: - DimVectorT contiguity; - DimVectorT divisibility; - DimVectorT constancy; - - // The constant value of the lattice if we can infer it. - std::optional constantValue; -}; - -class AxisInfoVisitor { -public: - AxisInfoVisitor() = default; - virtual ~AxisInfoVisitor() = default; - - bool isContiguousDim(const AxisInfo &info, ArrayRef shape, int dim) { - return info.getContiguity(dim) == shape[dim]; - } - - bool isConstantDim(const AxisInfo &info, ArrayRef shape, int dim) { - return info.getConstancy(dim) == shape[dim]; - } - - virtual AxisInfo - getAxisInfo(Operation *op, - ArrayRef *> operands) = 0; - - virtual bool match(Operation *op) = 0; -}; - -class AxisInfoVisitorList { -public: - template > - void append() { - (visitors.emplace_back(std::make_unique()), ...); - } - - AxisInfo apply(Operation *op, - ArrayRef *> operands) { - for (auto &visitor : visitors) - if (visitor->match(op)) - return visitor->getAxisInfo(op, operands); - return AxisInfo(); - } - -private: - std::vector> visitors; -}; - -namespace axisinfo { -using CallbackType = std::function; -} // namespace axisinfo - -// Module level axis info analysis based on the call graph, assuming that we do -// not have recursive functions. -// -// Since each function will be called multiple times, we need to calculate the -// axis info based on the axis info of all the callers. In the future, we can -// perform optimization using function cloning so that each call site will have -// unique axis info. -using AxisInfoMapT = DenseMap; -class ModuleAxisInfoAnalysis : public CallGraph { -public: - explicit ModuleAxisInfoAnalysis(ModuleOp moduleOp, - axisinfo::CallbackType callback = nullptr) - : CallGraph(moduleOp) { - SmallVector funcs; - walk( - // Pre-order edge walk callback - [](CallOpInterface callOp, FunctionOpInterface funcOp) {}, - // Post-order node walk callback - [&](FunctionOpInterface funcOp) { - funcs.push_back(funcOp); - funcMap.try_emplace(funcOp, AxisInfoMapT{}); - }); - SetVector sortedFuncs(funcs.begin(), funcs.end()); - SymbolTableCollection symbolTable; - for (auto funcOp : llvm::reverse(sortedFuncs)) { - initialize(funcOp, callback); - funcOp.walk([&](CallOpInterface callOp) { - auto callee = dyn_cast( - callOp.resolveCallableInTable(&symbolTable)); - update(callOp, callee); - }); - } - } - - AxisInfo *getAxisInfo(Value value) { - auto funcOp = - value.getParentRegion()->getParentOfType(); - auto *axisInfoMap = getFuncData(funcOp); - if (!axisInfoMap) { - return nullptr; - } - auto it = axisInfoMap->find(value); - if (it == axisInfoMap->end()) { - return nullptr; - } - return &(it->second); - } - - unsigned getContiguity(Value value); - unsigned getAlignment(Value value); - - // Overloads of the above methods but have separated elementBitWidth to - // calculate the contiguity. These are useful for computing axis info when - // lowering to hardware intrinsics that require a scalar/warp-uniform base ptr - // with separate per lane offsets like AMD buffer operations. - // - // As a concrete example, instead of a single tensor<128x64x!tt.ptr> - // value, now we have two separate values: !tt.ptr for the base pointer - // and tensor<128x64xi32> for the offset. For such cases, we want to compute - // the contiguity on the offsets but use the pointee element type bit width - // instead of the offset element type bit width for alignment - unsigned getContiguity(Value offsetsValue, unsigned elementBitWidth); - unsigned getAlignment(Value offsetsValue, unsigned elementBitWidth); - - unsigned getMaskAlignment(Value mask); - -private: - void initialize(FunctionOpInterface funcOp, - axisinfo::CallbackType callback = nullptr); - void update(CallOpInterface callOp, FunctionOpInterface funcOp); -}; -} // namespace mlir::triton - -#endif diff --git a/third_party/mthreads/include/triton/CMakeLists.txt b/third_party/mthreads/include/triton/CMakeLists.txt deleted file mode 100644 index 27c703b3cf..0000000000 --- a/third_party/mthreads/include/triton/CMakeLists.txt +++ /dev/null @@ -1,3 +0,0 @@ -add_subdirectory(Conversion) -add_subdirectory(Dialect) -add_subdirectory(Target) diff --git a/third_party/mthreads/include/triton/Conversion/CMakeLists.txt b/third_party/mthreads/include/triton/Conversion/CMakeLists.txt deleted file mode 100644 index 730f5cadd2..0000000000 --- a/third_party/mthreads/include/triton/Conversion/CMakeLists.txt +++ /dev/null @@ -1,2 +0,0 @@ -add_subdirectory(TritonGPUToLLVM) -add_subdirectory(TritonToTritonGPU) diff --git a/third_party/mthreads/include/triton/Conversion/MLIRTypes.h b/third_party/mthreads/include/triton/Conversion/MLIRTypes.h deleted file mode 100644 index dd8d4be4c2..0000000000 --- a/third_party/mthreads/include/triton/Conversion/MLIRTypes.h +++ /dev/null @@ -1,46 +0,0 @@ -#ifndef TRITON_CONVERSION_MLIR_TYPES_H -#define TRITON_CONVERSION_MLIR_TYPES_H - -#include "mlir/Transforms/DialectConversion.h" -#include "triton/Dialect/TritonGPU/IR/Dialect.h" - -// This file redefines some common MLIR types for easy usage. -namespace mlir { -namespace triton { -namespace type { - -// Integer types -inline Type i32Ty(MLIRContext *ctx) { return IntegerType::get(ctx, 32); } -inline Type i16Ty(MLIRContext *ctx) { return IntegerType::get(ctx, 16); } -inline Type i8Ty(MLIRContext *ctx) { return IntegerType::get(ctx, 8); } -inline Type u32Ty(MLIRContext *ctx) { - return IntegerType::get(ctx, 32, IntegerType::Unsigned); -} -inline Type u1Ty(MLIRContext *ctx) { - return IntegerType::get(ctx, 1, IntegerType::Unsigned); -} - -// Float types -inline Type f16Ty(MLIRContext *ctx) { return Float16Type::get(ctx); } -inline Type f32Ty(MLIRContext *ctx) { return Float32Type::get(ctx); } -inline Type f64Ty(MLIRContext *ctx) { return Float64Type::get(ctx); } -inline Type bf16Ty(MLIRContext *ctx) { return BFloat16Type::get(ctx); } - -inline bool isFloat8(Type type) { - return isa(type); -} - -inline bool isFloat(Type type) { - return type.isF32() || type.isF64() || type.isF16() || type.isF128() || - type.isBF16() || llvm::isa(type) || - isFloat8(type); -} - -inline bool isInt(Type type) { return type.isIntOrFloat() && !isFloat(type); } - -} // namespace type -} // namespace triton -} // namespace mlir - -#endif // TRITON_CONVERSION_MLIR_TYPES_H diff --git a/third_party/mthreads/include/triton/Conversion/TritonGPUToLLVM/AllocateSharedMemoryUtility.h b/third_party/mthreads/include/triton/Conversion/TritonGPUToLLVM/AllocateSharedMemoryUtility.h deleted file mode 100644 index 46a06ac65d..0000000000 --- a/third_party/mthreads/include/triton/Conversion/TritonGPUToLLVM/AllocateSharedMemoryUtility.h +++ /dev/null @@ -1,17 +0,0 @@ -#ifndef TRITON_CONVERSION_TRITON_GPU_TO_LLVM_ALLOCATE_UTILITY_H_ -#define TRITON_CONVERSION_TRITON_GPU_TO_LLVM_ALLOCATE_UTILITY_H_ - -#include "mlir/IR/BuiltinOps.h" -#include "triton/Analysis/Allocation.h" - -namespace mlir::triton::gpu { - -/// Attach shared memory related attributes to module and operations inside it. -/// This includes total shared memory consumption in module and shared memory -/// offsets of buffers associated with operations. -void attachAllocationSizeAndOffsetAttr(ModuleOp mod, - ModuleAllocation &allocation); - -} // namespace mlir::triton::gpu - -#endif // TRITON_CONVERSION_TRITON_GPU_TO_LLVM_ALLOCATE_UTILITY_H_ diff --git a/third_party/mthreads/include/triton/Conversion/TritonGPUToLLVM/AsmFormat.h b/third_party/mthreads/include/triton/Conversion/TritonGPUToLLVM/AsmFormat.h deleted file mode 100644 index 00ec880890..0000000000 --- a/third_party/mthreads/include/triton/Conversion/TritonGPUToLLVM/AsmFormat.h +++ /dev/null @@ -1,27 +0,0 @@ -#ifndef TRITON_CONVERSION_TRITON_GPU_TO_LLVM_ASM_FORMAT_H_ -#define TRITON_CONVERSION_TRITON_GPU_TO_LLVM_ASM_FORMAT_H_ - -#include "mlir/IR/Value.h" -#include "triton/Dialect/Triton/IR/Dialect.h" -#include "llvm/ADT/SmallVector.h" -#include "llvm/ADT/StringExtras.h" -#include "llvm/ADT/StringRef.h" -#include -#include - -namespace mlir { -class ConversionPatternRewriter; -class Location; - -namespace triton { -using llvm::StringRef; - -inline std::string strJoin(llvm::ArrayRef strs, - llvm::StringRef delimiter) { - return llvm::join(strs.begin(), strs.end(), delimiter); -} - -} // namespace triton -} // namespace mlir - -#endif // TRITON_CONVERSION_TRITON_GPU_TO_LLVM_ASM_FORMAT_H_ diff --git a/third_party/mthreads/include/triton/Conversion/TritonGPUToLLVM/CMakeLists.txt b/third_party/mthreads/include/triton/Conversion/TritonGPUToLLVM/CMakeLists.txt deleted file mode 100644 index 93f8374e59..0000000000 --- a/third_party/mthreads/include/triton/Conversion/TritonGPUToLLVM/CMakeLists.txt +++ /dev/null @@ -1,3 +0,0 @@ -set(LLVM_TARGET_DEFINITIONS Passes.td) -mlir_tablegen(Passes.h.inc -gen-pass-decls --name TritonGPUToLLVM) -add_public_tablegen_target(TritonGPUConversionPassIncGen) diff --git a/third_party/mthreads/include/triton/Conversion/TritonGPUToLLVM/FMADotUtility.h b/third_party/mthreads/include/triton/Conversion/TritonGPUToLLVM/FMADotUtility.h deleted file mode 100644 index 907d36ed45..0000000000 --- a/third_party/mthreads/include/triton/Conversion/TritonGPUToLLVM/FMADotUtility.h +++ /dev/null @@ -1,35 +0,0 @@ -#ifndef TRITON_CONVERSION_FMA_DOT_UTILITY_H -#define TRITON_CONVERSION_FMA_DOT_UTILITY_H - -#include "mlir/Conversion/LLVMCommon/TypeConverter.h" -#include "mlir/Support/LLVM.h" -#include "mlir/Transforms/DialectConversion.h" -#include "triton/Dialect/Triton/IR/Dialect.h" - -namespace mlir::triton::gpu { - -/// Abstract interface for scalar multiplication of Value vectors. -/// -/// Enable generation of hardware specific code in different backends. -class FMAVectorMultiplier { -public: - /// \returns scalar product of two arrays, plus c: a·b + c - virtual Value multiplyVectors(ArrayRef a, ArrayRef b, - Value c) = 0; - - virtual ~FMAVectorMultiplier() = default; -}; - -/// Implements a framework for FMA dot conversion to llvm. -/// -/// This function implements architecture independent part of FMA dot -/// conversion and calls "multiplier" object, which is defined by caller -/// and implements architecture dependant part of conversion. -LogicalResult parametricConvertFMADot(DotOp op, DotOp::Adaptor adaptor, - const LLVMTypeConverter *typeConverter, - ConversionPatternRewriter &rewriter, - FMAVectorMultiplier &multiplier); - -} // namespace mlir::triton::gpu - -#endif // TRITON_CONVERSION_FMA_DOT_UTILITY_H diff --git a/third_party/mthreads/include/triton/Conversion/TritonGPUToLLVM/Passes.h b/third_party/mthreads/include/triton/Conversion/TritonGPUToLLVM/Passes.h deleted file mode 100644 index 2a3a67a594..0000000000 --- a/third_party/mthreads/include/triton/Conversion/TritonGPUToLLVM/Passes.h +++ /dev/null @@ -1,25 +0,0 @@ -#ifndef TRITONGPU_CONVERSION_TRITONGPUTOLLVM_PASSES_H -#define TRITONGPU_CONVERSION_TRITONGPUTOLLVM_PASSES_H - -#include "mlir/Pass/Pass.h" - -#include - -namespace mlir { - -class ModuleOp; -template class OperationPass; - -namespace triton::gpu { - -#define GEN_PASS_DECL -#include "triton/Conversion/TritonGPUToLLVM/Passes.h.inc" - -#define GEN_PASS_REGISTRATION -#include "triton/Conversion/TritonGPUToLLVM/Passes.h.inc" - -} // namespace triton::gpu - -} // namespace mlir - -#endif diff --git a/third_party/mthreads/include/triton/Conversion/TritonGPUToLLVM/Passes.td b/third_party/mthreads/include/triton/Conversion/TritonGPUToLLVM/Passes.td deleted file mode 100644 index fa3cc63c72..0000000000 --- a/third_party/mthreads/include/triton/Conversion/TritonGPUToLLVM/Passes.td +++ /dev/null @@ -1,45 +0,0 @@ -#ifndef TRITONCOMMONGPU_CONVERSION_PASSES -#define TRITONCOMMONGPU_CONVERSION_PASSES - -include "mlir/Pass/PassBase.td" - -def AllocateSharedMemory : Pass<"allocate-shared-memory", "mlir::ModuleOp"> { - let summary = "Add metadata for shared memory allocation"; - - let description = [{ - This pass uses the `ModuleAllocation` analysis to: - - Annotate modules with an attribute with the amount of shared/local - memory used. - - Annotate operations with an offset into the total shared/local memory. - }]; -} - -def TritonGPUGlobalScratchAllocationPass : Pass<"tritongpu-global-scratch-memory-allocation", "mlir::ModuleOp"> { - let summary = "Assign global scratch memory allocation"; - - let description = [{ - Decide on global scratch space memory allocation and assign attributes to each allocation. - }]; - - let dependentDialects = [ - "mlir::triton::gpu::TritonGPUDialect" - ]; -} - -def TritonGPUAllocateWarpGroups : Pass<"tritongpu-allocate-warp-groups", "mlir::ModuleOp"> { - let summary = "Allocate warp groups"; - - let description = [{ - The `tritongpu-allocate-warp-groups` pass performs warpgroup allocation for - a GPU program. When a GPU program contains warp specialization, additional - warps are launched in addition to the "default" warp group. The "default" - warpgroup executes top-level code in a `tt.func` and its size is specified - by the user via the `num_warps` argument. - - This pass analyzes `ttg.warp_specialize` ops in the program and determines - the total number of needed warps, then attaches the range of warp IDs to - each warpgroup function. - }]; -} - -#endif diff --git a/third_party/mthreads/include/triton/Conversion/TritonGPUToLLVM/TypeConverter.h b/third_party/mthreads/include/triton/Conversion/TritonGPUToLLVM/TypeConverter.h deleted file mode 100644 index 1adbbee4e3..0000000000 --- a/third_party/mthreads/include/triton/Conversion/TritonGPUToLLVM/TypeConverter.h +++ /dev/null @@ -1,39 +0,0 @@ -#ifndef TRITON_CONVERSION_TRITONGPU_TO_LLVM_TYPECONVERTER_H -#define TRITON_CONVERSION_TRITONGPU_TO_LLVM_TYPECONVERTER_H - -#include "mlir/Conversion/LLVMCommon/TypeConverter.h" -#include "mlir/Dialect/LLVMIR/LLVMDialect.h" -#include "triton/Conversion/MLIRTypes.h" -#include "triton/Conversion/TritonGPUToLLVM/TargetInfoBase.h" -#include "triton/Dialect/TritonGPU/IR/Types.h" - -using namespace mlir; -using namespace mlir::triton; - -class TritonGPUToLLVMTypeConverter : public LLVMTypeConverter { -public: - using TypeConverter::convertType; - - TritonGPUToLLVMTypeConverter(MLIRContext *ctx, - const LowerToLLVMOptions &option, - const TargetInfoBase &targetInfo, - const DataLayoutAnalysis *analysis = nullptr); - TritonGPUToLLVMTypeConverter(MLIRContext *ctx, - const TargetInfoBase &targetInfo, - const DataLayoutAnalysis *analysis = nullptr); - - Type convertTritonTensorType(RankedTensorType type, - const TargetInfoBase &targetInfo); - Type convertMemDescType(triton::gpu::MemDescType type, - const TargetInfoBase &targetInfo); - Type convertAsyncTokenType(triton::gpu::AsyncTokenType type); - - template void convertFP8Type() { - (addConversion([&](T type) -> std::optional { - return IntegerType::get(type.getContext(), 8); - }), - ...); - } -}; - -#endif diff --git a/third_party/mthreads/include/triton/Conversion/TritonToTritonGPU/CMakeLists.txt b/third_party/mthreads/include/triton/Conversion/TritonToTritonGPU/CMakeLists.txt deleted file mode 100644 index 99d90c4d75..0000000000 --- a/third_party/mthreads/include/triton/Conversion/TritonToTritonGPU/CMakeLists.txt +++ /dev/null @@ -1,3 +0,0 @@ -set(LLVM_TARGET_DEFINITIONS Passes.td) -mlir_tablegen(Passes.h.inc -gen-pass-decls --name TritonToTritonGPU) -add_public_tablegen_target(TritonConversionPassIncGen) diff --git a/third_party/mthreads/include/triton/Conversion/TritonToTritonGPU/Passes.h b/third_party/mthreads/include/triton/Conversion/TritonToTritonGPU/Passes.h deleted file mode 100644 index 054f9ea959..0000000000 --- a/third_party/mthreads/include/triton/Conversion/TritonToTritonGPU/Passes.h +++ /dev/null @@ -1,15 +0,0 @@ -#ifndef TRITON_CONVERSION_PASSES_H -#define TRITON_CONVERSION_PASSES_H - -#include "mlir/Pass/Pass.h" - -namespace mlir::triton { - -#define GEN_PASS_DECL -#include "triton/Conversion/TritonToTritonGPU/Passes.h.inc" -#define GEN_PASS_REGISTRATION -#include "triton/Conversion/TritonToTritonGPU/Passes.h.inc" - -} // namespace mlir::triton - -#endif diff --git a/third_party/mthreads/include/triton/Conversion/TritonToTritonGPU/Passes.td b/third_party/mthreads/include/triton/Conversion/TritonToTritonGPU/Passes.td deleted file mode 100644 index 2449637eb1..0000000000 --- a/third_party/mthreads/include/triton/Conversion/TritonToTritonGPU/Passes.td +++ /dev/null @@ -1,56 +0,0 @@ -#ifndef TRITON_CONVERSION_PASSES -#define TRITON_CONVERSION_PASSES - -include "mlir/Pass/PassBase.td" - -def ConvertTritonToTritonGPU: Pass<"convert-triton-to-tritongpu", "mlir::ModuleOp"> { - let summary = "Convert Triton to TritonGPU"; - let description = [{ - This pass converts the Triton Dialect into the TritonGPU Dialect. - This is a partial conversion that also affects other dialects - (namely `Arith`, `Math`, `SCF` and `CF`). - For these dialects, and many Triton dialect operations the conversions - mainly consists of enhancing the tensor type and the `tt.ptr>` - type with an appropriate layout encoding (these encodings generally - include information on `numWarps`, `threadsPerWarp` and `numCTAs`). - }]; - - let dependentDialects = ["mlir::arith::ArithDialect", - "mlir::math::MathDialect", - // TODO: Does this pass depend on SCF? - "mlir::scf::SCFDialect", - "mlir::triton::TritonDialect", - "mlir::triton::gpu::TritonGPUDialect"]; - - let options = [ - Option<"target", "target", - "std::string", /*default*/"\"\"", - "the GPU target, e.g., cuda:80, hip:gfx942">, - Option<"numWarps", "num-warps", - "int32_t", /*default*/"4", - "number of warps">, - Option<"threadsPerWarp", "threads-per-warp", - "int32_t", /*default*/"32", - "number of threads per warp">, - Option<"numCTAs", "num-ctas", - "int32_t", /*default*/"1", - "number of ctas in a cga">, - Option<"enableSourceRemat", "enable-source-remat", - "bool", /*default*/"false", - "enable trivial source rematerialization">, - ]; -} - -def RelayoutTritonGPU : Pass<"relayout-tritongpu", "mlir::ModuleOp"> { - let summary = "relayout pass for `ttg` and `ttng` operations"; - let description = [{ - The `relayout-tritongpu` pass is used during relayout of TTGIR - during warp specialization. Warp specialization may change the number of - warps for a partition, which requires reassigning layouts to all the - operations in the partition. However, those operations may include TritonGPU - and TritonNvidiaGPU dialect operations with specific layout requirements, - so they have to be re-inferred during this pass. - }]; -} - -#endif diff --git a/third_party/mthreads/include/triton/Dialect/CMakeLists.txt b/third_party/mthreads/include/triton/Dialect/CMakeLists.txt deleted file mode 100644 index 19ca22ec3b..0000000000 --- a/third_party/mthreads/include/triton/Dialect/CMakeLists.txt +++ /dev/null @@ -1,7 +0,0 @@ -add_subdirectory(Triton) -add_subdirectory(TritonGPU) -add_subdirectory(TritonNvidiaGPU) -add_subdirectory(TritonInstrument) -add_subdirectory(Gluon) -add_subdirectory(NVGPU) -add_subdirectory(NVWS) diff --git a/third_party/mthreads/include/triton/Dialect/Gluon/CMakeCache.txt b/third_party/mthreads/include/triton/Dialect/Gluon/CMakeCache.txt deleted file mode 100644 index 9f57627c32..0000000000 --- a/third_party/mthreads/include/triton/Dialect/Gluon/CMakeCache.txt +++ /dev/null @@ -1,2 +0,0 @@ -add_subdirectory(IR) -add_subdirectory(Transforms) diff --git a/third_party/mthreads/include/triton/Dialect/Gluon/CMakeLists.txt b/third_party/mthreads/include/triton/Dialect/Gluon/CMakeLists.txt deleted file mode 100644 index 9f57627c32..0000000000 --- a/third_party/mthreads/include/triton/Dialect/Gluon/CMakeLists.txt +++ /dev/null @@ -1,2 +0,0 @@ -add_subdirectory(IR) -add_subdirectory(Transforms) diff --git a/third_party/mthreads/include/triton/Dialect/Gluon/IR/CMakeLists.txt b/third_party/mthreads/include/triton/Dialect/Gluon/IR/CMakeLists.txt deleted file mode 100644 index 8e42fc0904..0000000000 --- a/third_party/mthreads/include/triton/Dialect/Gluon/IR/CMakeLists.txt +++ /dev/null @@ -1,17 +0,0 @@ -set(MLIR_BINARY_DIR ${CMAKE_BINARY_DIR}) - -set(LLVM_TARGET_DEFINITIONS GluonOps.td) -mlir_tablegen(Ops.h.inc -gen-op-decls) -mlir_tablegen(Ops.cpp.inc -gen-op-defs) -add_mlir_doc(GluonOps GluonOps dialects/ -gen-op-doc) - -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) -add_mlir_doc(GluonDialect GluonDialect dialects/ -gen-dialect-doc) - -set(LLVM_TARGET_DEFINITIONS GluonAttrDefs.td) -mlir_tablegen(GluonAttrDefs.h.inc -gen-attrdef-decls) -mlir_tablegen(GluonAttrDefs.cpp.inc -gen-attrdef-defs) - -add_public_tablegen_target(GluonTableGen) diff --git a/third_party/mthreads/include/triton/Dialect/Gluon/IR/Dialect.h b/third_party/mthreads/include/triton/Dialect/Gluon/IR/Dialect.h deleted file mode 100644 index 3004e71a62..0000000000 --- a/third_party/mthreads/include/triton/Dialect/Gluon/IR/Dialect.h +++ /dev/null @@ -1,11 +0,0 @@ -#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" - -#define GET_OP_CLASSES -#include "triton/Dialect/Gluon/IR/Ops.h.inc" diff --git a/third_party/mthreads/include/triton/Dialect/Gluon/IR/GluonAttrDefs.td b/third_party/mthreads/include/triton/Dialect/Gluon/IR/GluonAttrDefs.td deleted file mode 100644 index f2b0da23a9..0000000000 --- a/third_party/mthreads/include/triton/Dialect/Gluon/IR/GluonAttrDefs.td +++ /dev/null @@ -1,23 +0,0 @@ -#ifndef GLUON_ATTRDEFS -#define GLUON_ATTRDEFS - -include "mlir/IR/AttrTypeBase.td" -include "triton/Dialect/Gluon/IR/GluonDialect.td" - -def Gluon_AutoEncodingAttr : AttrDef { - let mnemonic = "auto_encoding"; - let attrName = "gluon.auto_encoding"; - let description = [{ - An encoding that is inferred from neighboring ops in the graph. - }]; -} - -def Gluon_CoalescedEncodingAttr : AttrDef { - let mnemonic = "coalesced_encoding"; - let attrName = "gluon.coalesced_encoding"; - let description = [{ - An encoding that is optimized for load/store performance. - }]; -} - -#endif diff --git a/third_party/mthreads/include/triton/Dialect/Gluon/IR/GluonDialect.td b/third_party/mthreads/include/triton/Dialect/Gluon/IR/GluonDialect.td deleted file mode 100644 index 37e55f12ed..0000000000 --- a/third_party/mthreads/include/triton/Dialect/Gluon/IR/GluonDialect.td +++ /dev/null @@ -1,22 +0,0 @@ -#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 diff --git a/third_party/mthreads/include/triton/Dialect/Gluon/IR/GluonOps.td b/third_party/mthreads/include/triton/Dialect/Gluon/IR/GluonOps.td deleted file mode 100644 index d268c0e515..0000000000 --- a/third_party/mthreads/include/triton/Dialect/Gluon/IR/GluonOps.td +++ /dev/null @@ -1,32 +0,0 @@ -#ifndef GLUON_OPS -#define GLUON_OPS - -include "triton/Dialect/Gluon/IR/GluonDialect.td" -include "triton/Dialect/Gluon/IR/GluonAttrDefs.td" -include "triton/Dialect/Triton/IR/TritonInterfaces.td" -include "triton/Dialect/Triton/IR/TritonTypes.td" - -class Gluon_Op traits = []> : - Op { -} - -def Gluon_SetAutoLayoutOp : Gluon_Op<"set_auto_layout", - [SameOperandsAndResultShape, - SameOperandsAndResultElementType]> { - let summary = "set auto encoding to a concrete encoding type"; - - let arguments = (ins TT_Tensor:$src); - - let results = (outs TT_Tensor:$result); - - let builders = [ - OpBuilder<(ins "Attribute":$encoding, "Value":$value)> - ]; - - let hasVerifier = 1; - - let assemblyFormat = "$src attr-dict `:` type($src) `->` type($result)"; -} - -#endif // GLUON_OPS diff --git a/third_party/mthreads/include/triton/Dialect/Gluon/Transforms/CMakeLists.txt b/third_party/mthreads/include/triton/Dialect/Gluon/Transforms/CMakeLists.txt deleted file mode 100644 index a2d298d0c1..0000000000 --- a/third_party/mthreads/include/triton/Dialect/Gluon/Transforms/CMakeLists.txt +++ /dev/null @@ -1,3 +0,0 @@ -set(LLVM_TARGET_DEFINITIONS Passes.td) -mlir_tablegen(Passes.h.inc -gen-pass-decls -name Gluon) -add_public_tablegen_target(GluonTransformsIncGen) diff --git a/third_party/mthreads/include/triton/Dialect/Gluon/Transforms/InferLayoutUtils.h b/third_party/mthreads/include/triton/Dialect/Gluon/Transforms/InferLayoutUtils.h deleted file mode 100644 index 3cd4b0d508..0000000000 --- a/third_party/mthreads/include/triton/Dialect/Gluon/Transforms/InferLayoutUtils.h +++ /dev/null @@ -1,20 +0,0 @@ -#ifndef TRITON_DIALECT_GLUON_TRANSFORMS_INFERLAYOUTUTILS_H_ -#define TRITON_DIALECT_GLUON_TRANSFORMS_INFERLAYOUTUTILS_H_ - -#include "triton/Dialect/Gluon/IR/Dialect.h" -#include "triton/Dialect/Gluon/Transforms/Passes.h" -#include "llvm/ADT/MapVector.h" -#include "llvm/ADT/PriorityWorklist.h" - -namespace mlir::triton::gluon { - -LogicalResult -inferLayout(FuncOp func, llvm::function_ref typeCheck, - const SmallVector> &seedEncodings); - -LogicalResult doubleCheckEncodings(ModuleOp &mod, - llvm::function_ref typeCheck); - -} // namespace mlir::triton::gluon - -#endif // TRITON_DIALECT_GLUON_TRANSFORMS_INFERLAYOUTUTILS_H_ diff --git a/third_party/mthreads/include/triton/Dialect/Gluon/Transforms/Passes.h b/third_party/mthreads/include/triton/Dialect/Gluon/Transforms/Passes.h deleted file mode 100644 index 353d21e04f..0000000000 --- a/third_party/mthreads/include/triton/Dialect/Gluon/Transforms/Passes.h +++ /dev/null @@ -1,13 +0,0 @@ -#pragma once -#include "mlir/IR/BuiltinOps.h" -#include "mlir/Pass/Pass.h" -#include "triton/Dialect/Gluon/IR/Dialect.h" -#include - -namespace mlir::triton::gluon { - -#define GEN_PASS_DECL -#define GEN_PASS_REGISTRATION -#include "triton/Dialect/Gluon/Transforms/Passes.h.inc" - -} // namespace mlir::triton::gluon diff --git a/third_party/mthreads/include/triton/Dialect/NVGPU/CMakeLists.txt b/third_party/mthreads/include/triton/Dialect/NVGPU/CMakeLists.txt deleted file mode 100644 index f33061b2d8..0000000000 --- a/third_party/mthreads/include/triton/Dialect/NVGPU/CMakeLists.txt +++ /dev/null @@ -1 +0,0 @@ -add_subdirectory(IR) diff --git a/third_party/mthreads/include/triton/Dialect/NVWS/CMakeLists.txt b/third_party/mthreads/include/triton/Dialect/NVWS/CMakeLists.txt deleted file mode 100644 index 9f57627c32..0000000000 --- a/third_party/mthreads/include/triton/Dialect/NVWS/CMakeLists.txt +++ /dev/null @@ -1,2 +0,0 @@ -add_subdirectory(IR) -add_subdirectory(Transforms) diff --git a/third_party/mthreads/include/triton/Dialect/Triton/CMakeLists.txt b/third_party/mthreads/include/triton/Dialect/Triton/CMakeLists.txt deleted file mode 100644 index 9f57627c32..0000000000 --- a/third_party/mthreads/include/triton/Dialect/Triton/CMakeLists.txt +++ /dev/null @@ -1,2 +0,0 @@ -add_subdirectory(IR) -add_subdirectory(Transforms) diff --git a/third_party/mthreads/include/triton/Dialect/Triton/IR/CMakeLists.txt b/third_party/mthreads/include/triton/Dialect/Triton/IR/CMakeLists.txt deleted file mode 100644 index 5ed50b3a23..0000000000 --- a/third_party/mthreads/include/triton/Dialect/Triton/IR/CMakeLists.txt +++ /dev/null @@ -1,31 +0,0 @@ -set(MLIR_BINARY_DIR ${CMAKE_BINARY_DIR}) - -set(LLVM_TARGET_DEFINITIONS TritonOps.td) -mlir_tablegen(Ops.h.inc -gen-op-decls) -mlir_tablegen(Ops.cpp.inc -gen-op-defs) -mlir_tablegen(OpsEnums.h.inc -gen-enum-decls) -mlir_tablegen(OpsEnums.cpp.inc -gen-enum-defs) -add_mlir_doc(TritonOps TritonOps dialects/ -gen-op-doc) - -set(LLVM_TARGET_DEFINITIONS TritonDialect.td) -mlir_tablegen(Dialect.h.inc -gen-dialect-decls) -mlir_tablegen(Dialect.cpp.inc -gen-dialect-defs) -add_mlir_doc(TritonDialect TritonDialect dialects/ -gen-dialect-doc) - -set(LLVM_TARGET_DEFINITIONS TritonTypes.td) -mlir_tablegen(Types.h.inc -gen-typedef-decls) -mlir_tablegen(Types.cpp.inc -gen-typedef-defs) - -set(LLVM_TARGET_DEFINITIONS TritonInterfaces.td) -mlir_tablegen(AttrInterfaces.h.inc -gen-attr-interface-decls) -mlir_tablegen(AttrInterfaces.cpp.inc -gen-attr-interface-defs) - -set(LLVM_TARGET_DEFINITIONS TritonOpInterfaces.td) -mlir_tablegen(OpInterfaces.h.inc -gen-op-interface-decls) -mlir_tablegen(OpInterfaces.cpp.inc -gen-op-interface-defs) - -set(LLVM_TARGET_DEFINITIONS TritonTypeInterfaces.td) -mlir_tablegen(TypeInterfaces.h.inc -gen-type-interface-decls) -mlir_tablegen(TypeInterfaces.cpp.inc -gen-type-interface-defs) - -add_public_tablegen_target(TritonTableGen) diff --git a/third_party/mthreads/include/triton/Dialect/Triton/IR/DiscardableAttributes.h b/third_party/mthreads/include/triton/Dialect/Triton/IR/DiscardableAttributes.h deleted file mode 100644 index 68908fa926..0000000000 --- a/third_party/mthreads/include/triton/Dialect/Triton/IR/DiscardableAttributes.h +++ /dev/null @@ -1,15 +0,0 @@ -#ifndef TRITON_DIALECT_TRITON_IR_DISCARDABLE_ATTRIBUTES_H_ -#define TRITON_DIALECT_TRITON_IR_DISCARDABLE_ATTRIBUTES_H_ - -#include "mlir/Support/LLVM.h" -#include "triton/Dialect/Triton/IR/Dialect.h" - -namespace mlir::triton { - -// Filter out attributes from the given operation that are not present in -// the allowList. -[[nodiscard]] SmallVector -filterDiscardableAttrs(Operation *op, ArrayRef allowList); - -} // namespace mlir::triton -#endif // TRITON_DIALECT_TRITON_IR_DISCARDABLE_ATTRIBUTES_H_ diff --git a/third_party/mthreads/include/triton/Dialect/Triton/IR/Interfaces.h b/third_party/mthreads/include/triton/Dialect/Triton/IR/Interfaces.h deleted file mode 100644 index fb5951fa5c..0000000000 --- a/third_party/mthreads/include/triton/Dialect/Triton/IR/Interfaces.h +++ /dev/null @@ -1,45 +0,0 @@ -#ifndef TRITON_IR_INTERFACES_H_ -#define TRITON_IR_INTERFACES_H_ - -#include "mlir/IR/DialectImplementation.h" -#include "mlir/IR/OpDefinition.h" -#include "mlir/Transforms/InliningUtils.h" - -#define GET_TYPEDEF_CLASSES -#include "triton/Dialect/Triton/IR/AttrInterfaces.h.inc" - -namespace mlir::triton { - -//===----------------------------------------------------------------------===// -// TritonDialect Dialect Interfaces -//===----------------------------------------------------------------------===// - -struct TritonInlinerInterface : public DialectInlinerInterface { - using DialectInlinerInterface::DialectInlinerInterface; - - bool isLegalToInline(Operation *call, Operation *callable, - bool wouldBeCloned) const final; - bool isLegalToInline(Region *dest, Region *src, bool wouldBeCloned, - IRMapping &valueMapping) const final { - return true; - } - bool isLegalToInline(Operation *, Region *, bool wouldBeCloned, - IRMapping &) const final { - return true; - } - - //===--------------------------------------------------------------------===// - // Transformation Hooks - //===--------------------------------------------------------------------===// - - /// Handle the given inlined terminator by replacing it with a new operation - /// as necessary. - void handleTerminator(Operation *op, Block *newDest) const final; - /// Handle the given inlined terminator by replacing it with a new operation - /// as necessary. - void handleTerminator(Operation *op, ValueRange valuesToRepl) const final; -}; - -} // namespace mlir::triton - -#endif // TRITON_IR_TYPES_H_ diff --git a/third_party/mthreads/include/triton/Dialect/Triton/IR/OpInterfaces.h b/third_party/mthreads/include/triton/Dialect/Triton/IR/OpInterfaces.h deleted file mode 100644 index 326f876e1c..0000000000 --- a/third_party/mthreads/include/triton/Dialect/Triton/IR/OpInterfaces.h +++ /dev/null @@ -1,24 +0,0 @@ -#ifndef TRITON_IR_OP_INTERFACES_H_ -#define TRITON_IR_OP_INTERFACES_H_ - -#include "mlir/IR/OpDefinition.h" -#include "triton/Dialect/Triton/IR/Types.h" - -namespace mlir { - -namespace triton { - -namespace impl { - -LogicalResult verifyTransposeOpInterface(Operation *op); - -LogicalResult verifyDotOpInterface(Operation *op); - -} // namespace impl - -} // namespace triton -} // namespace mlir - -#include "triton/Dialect/Triton/IR/OpInterfaces.h.inc" - -#endif // TRITON_IR_OP_INTERFACES_H_ diff --git a/third_party/mthreads/include/triton/Dialect/Triton/IR/TritonAttrDefs.td b/third_party/mthreads/include/triton/Dialect/Triton/IR/TritonAttrDefs.td deleted file mode 100644 index 5a76a1d7b1..0000000000 --- a/third_party/mthreads/include/triton/Dialect/Triton/IR/TritonAttrDefs.td +++ /dev/null @@ -1,154 +0,0 @@ -#ifndef TRITON_ATTR_DEFS -#define TRITON_ATTR_DEFS - -include "mlir/IR/EnumAttr.td" - -// Attributes for LoadOp and StoreOp -def TT_CacheModifierAttr : I32EnumAttr< - "CacheModifier", "", - [ - I32EnumAttrCase<"NONE", 1, "none">, - I32EnumAttrCase<"CA", 2, "ca">, - I32EnumAttrCase<"CG", 3, "cg">, - I32EnumAttrCase<"WB", 4, "wb">, - I32EnumAttrCase<"CS", 5, "cs">, - I32EnumAttrCase<"WT", 6, "wt">, - I32EnumAttrCase<"CV", 7, "cv">, - ]> { - let cppNamespace = "::mlir::triton"; -} - -def TT_MemSemanticAttr : I32EnumAttr< - "MemSemantic", "", - [ - I32EnumAttrCase<"RELAXED", 1, "relaxed">, - I32EnumAttrCase<"ACQUIRE", 2, "acquire">, - I32EnumAttrCase<"RELEASE", 3, "release">, - I32EnumAttrCase<"ACQUIRE_RELEASE", 4, "acq_rel">, - ]> { - let cppNamespace = "::mlir::triton"; -} - -def TT_EvictionPolicyAttr : I32EnumAttr< - "EvictionPolicy", "", - [ - I32EnumAttrCase<"NORMAL", 1, "evict_normal">, - I32EnumAttrCase<"EVICT_FIRST", 2, "evict_first">, - I32EnumAttrCase<"EVICT_LAST", 3, "evict_last"> - ]> { - let cppNamespace = "::mlir::triton"; -} - -def TT_PaddingOptionAttr : I32EnumAttr< - "PaddingOption", "", - [ - I32EnumAttrCase<"PAD_ZERO", 1, "zero">, - // We can not set the string value to "NAN" because it is a keyword in C++ - I32EnumAttrCase<"PAD_NAN", 2, "nan"> - ]> { - let cppNamespace = "::mlir::triton"; -} - -// atomic -def TT_AtomicRMWAttr : I32EnumAttr< - "RMWOp", "", - [ - I32EnumAttrCase<"AND", 1, "and">, - I32EnumAttrCase<"OR", 2, "or">, - I32EnumAttrCase<"XOR", 3, "xor">, - I32EnumAttrCase<"ADD", 4, "add">, - I32EnumAttrCase<"FADD", 5, "fadd">, - I32EnumAttrCase<"MAX", 6, "max">, - I32EnumAttrCase<"MIN", 7, "min">, - I32EnumAttrCase<"UMAX", 8, "umax">, - I32EnumAttrCase<"UMIN", 9, "umin">, - I32EnumAttrCase<"XCHG", 10, "exch"> - ]> { - let cppNamespace = "::mlir::triton"; -} - -def TT_DescriptorReduceKindAttr : I32EnumAttr< - "DescriptorReduceKind", "", - [ - I32EnumAttrCase<"ADD", 1, "add">, - I32EnumAttrCase<"MIN", 2, "min">, - I32EnumAttrCase<"MAX", 3, "max">, - I32EnumAttrCase<"INC", 4, "inc">, - I32EnumAttrCase<"DEC", 5, "dec">, - I32EnumAttrCase<"AND", 6, "and">, - I32EnumAttrCase<"OR", 7, "or">, - I32EnumAttrCase<"XOR", 8, "xor">, - ]> { - let cppNamespace = "::mlir::triton"; -} - -def TT_MemSyncScopeAttr : I32EnumAttr< - "MemSyncScope", "", - [ - I32EnumAttrCase<"GPU", 1, "gpu">, - I32EnumAttrCase<"CTA", 2, "cta">, - I32EnumAttrCase<"SYSTEM", 3, "sys">, - ]> { - let cppNamespace = "::mlir::triton"; -} - -// Program ID dimensions. -def TT_ProgramDim : I32EnumAttr< - "ProgramIDDim", "", - [ - I32EnumAttrCase<"X", 0, "x">, - I32EnumAttrCase<"Y", 1, "y">, - I32EnumAttrCase<"Z", 2, "z">, - ]> { - let cppNamespace = "::mlir::triton"; -} - -// Rounding mode. -def TT_RoundingModeAttr : I32EnumAttr< - "RoundingMode", "", - [ - I32EnumAttrCase<"RTZ", 0, "rtz">, - I32EnumAttrCase<"RTNE", 1, "rtne">, - ]> { - let cppNamespace = "::mlir::triton"; -} - -// PropagateNan. -def TT_PropagateNanAttr : I32EnumAttr< - "PropagateNan", "", - [ - I32EnumAttrCase<"NONE", 0, "none">, - I32EnumAttrCase<"ALL", 0xFFFF, "all">, - ]> { - let cppNamespace = "::mlir::triton"; -} - -// InputPrecision -def TT_InputPrecisionAttr : I32EnumAttr< - "InputPrecision", "", - [ - I32EnumAttrCase<"TF32", 0, "tf32">, - I32EnumAttrCase<"TF32x3", 1, "tf32x3">, - I32EnumAttrCase<"IEEE", 2, "ieee">, - I32EnumAttrCase<"BF16x3", 3, "bf16x3">, - I32EnumAttrCase<"BF16x6", 4, "bf16x6"> - ]>{ - let cppNamespace = "::mlir::triton"; -} - -// Type for ScaleDotElemType kind of floats. -def TT_ScaleDotElemTypeAttr : I32EnumAttr< - "ScaleDotElemType", "", - [ - I32EnumAttrCase<"E4M3", 0, "e4m3">, - I32EnumAttrCase<"E5M2", 1, "e5m2">, - I32EnumAttrCase<"E2M3", 2, "e2m3">, - I32EnumAttrCase<"E3M2", 3, "e3m2">, - I32EnumAttrCase<"E2M1", 4, "e2m1">, - I32EnumAttrCase<"BF16", 5, "bf16">, - I32EnumAttrCase<"FP16", 6, "fp16"> - ]>{ - let cppNamespace = "::mlir::triton"; -} - -#endif diff --git a/third_party/mthreads/include/triton/Dialect/Triton/IR/TritonDialect.td b/third_party/mthreads/include/triton/Dialect/Triton/IR/TritonDialect.td deleted file mode 100644 index d0e25946b5..0000000000 --- a/third_party/mthreads/include/triton/Dialect/Triton/IR/TritonDialect.td +++ /dev/null @@ -1,60 +0,0 @@ -#ifndef TRITON_DIALECT -#define TRITON_DIALECT - -include "mlir/IR/OpBase.td" - -def Triton_Dialect : Dialect { - let name = "tt"; - - let cppNamespace = "::mlir::triton"; - - let summary = "The Triton IR in MLIR"; - - let description = [{ - Triton Dialect. - - Dependent Dialects: - * Arith: - * addf, addi, andi, cmpf, cmpi, divf, fptosi, ... - * Math: - * exp, sin, cos, log, ... - * StructuredControlFlow: - * for, if, while, yield, condition - * ControlFlow: - * br, cond_br - }]; - - let dependentDialects = [ - "arith::ArithDialect", - "math::MathDialect", - "scf::SCFDialect", - "cf::ControlFlowDialect", - "ub::UBDialect" - ]; - - let extraClassDeclaration = [{ - void registerTypes(); - - static TritonDialect *getLoaded(MLIRContext *ctx) { - return ctx->getLoadedDialect(); - } - static TritonDialect *getLoaded(Operation *op) { - return getLoaded(op->getContext()); - } - }]; - - let discardableAttrs = (ins - "::mlir::IntegerAttr":$num_stages, - "::mlir::IntegerAttr":$latency, - "::mlir::IntegerAttr":$self_latency - ); - - let hasConstantMaterializer = 1; - let useDefaultTypePrinterParser = 1; - let usePropertiesForAttributes = 1; -} - -include "triton/Dialect/Triton/IR/TritonTypes.td" - - -#endif // TRITON_DIALECT diff --git a/third_party/mthreads/include/triton/Dialect/Triton/IR/TritonOpInterfaces.td b/third_party/mthreads/include/triton/Dialect/Triton/IR/TritonOpInterfaces.td deleted file mode 100644 index 5cb7f8f333..0000000000 --- a/third_party/mthreads/include/triton/Dialect/Triton/IR/TritonOpInterfaces.td +++ /dev/null @@ -1,118 +0,0 @@ -#ifndef TRITON_OP_INTERFACES -#define TRITON_OP_INTERFACES - -include "mlir/IR/OpBase.td" - - -def TransposeOpInterface : OpInterface<"TransposeOpInterface"> { - let description = [{ - This interface is implemented by operations that perform a transpose. - It provides methods to access common properties such as the order attribute - and the source operand. - }]; - - let cppNamespace = "::mlir::triton"; - - let methods = [ - InterfaceMethod< - /*desc=*/"Get the source operand of the transposition.", - /*retType=*/"::mlir::Value", - /*methodName=*/"getSrc", - /*args=*/(ins)>, - InterfaceMethod< - /*desc=*/"Get the order of the transposition.", - /*retType=*/"::mlir::ArrayRef", - /*methodName=*/"getOrder", - /*args=*/(ins)> - ]; - - let verify = [{ - return ::mlir::triton::impl::verifyTransposeOpInterface($_op); - }]; -} - -def DotOpInterface : OpInterface<"DotOpInterface"> { - let description = [{ - This interface is implemented by operations that perform a dot product. - }]; - - let cppNamespace = "::mlir::triton"; - - let methods = [ - InterfaceMethod< - /*desc=*/"Get the LHS A tensor", - /*retType=*/"::mlir::Value", - /*methodName=*/"getA", - /*args=*/(ins)>, - InterfaceMethod< - /*desc=*/"Get the RHS B tensor", - /*retType=*/"::mlir::Value", - /*methodName=*/"getB", - /*args=*/(ins)>, - InterfaceMethod< - /*desc=*/"Get the output tensor", - /*retType=*/"::mlir::Value", - /*methodName=*/"getD", - /*args=*/(ins)>, - InterfaceMethod< - /*desc=*/"Verify the dimensions of the A and B DotOp operands.", - /*retType=*/"bool", - /*methodName=*/"verifyDims", - /*args=*/(ins)>, - InterfaceMethod< - /*desc=*/"Verify the dimensions of the DotOp output.", - /*retType=*/"bool", - /*methodName=*/"verifyOutputDims", - /*args=*/(ins), - /*methodBody=*/[{}], - /*defaultImpl=*/ [{ - auto aTy = cast($_op.getA().getType()); - auto bTy = cast($_op.getB().getType()); - auto cTy = cast($_op->getOperand(2).getType()); - auto dTy = cast($_op.getD().getType()); - auto aShape = aTy.getShape(); - auto bShape = bTy.getShape(); - auto cShape = cTy.getShape(); - return cShape[cShape.size() - 2] == aShape[aShape.size() - 2] && - cShape[cShape.size() - 1] == bShape[aShape.size() - 1]; - }]> - ]; - - let verify = [{ return ::mlir::triton::impl::verifyDotOpInterface($_op); }]; -} - -def TT_DescriptorOpInterface : OpInterface<"DescriptorOpInterface"> { - let description = [{ - Common interface to get the descriptor argument from an operation on tensor descriptors. - }]; - - let cppNamespace = "::mlir::triton"; - - let methods = [ - InterfaceMethod< - /*desc=*/"Get the descriptor", - /*retType=*/"::mlir::TypedValue", - /*methodName=*/"getDesc", - /*args=*/(ins)>, - ]; -} - -def TT_DescriptorStoreLikeOpInterface : OpInterface<"DescriptorStoreLikeOpInterface", [TT_DescriptorOpInterface]> { - let cppNamespace = "::mlir::triton"; - - let methods = [ - InterfaceMethod< - /*desc=*/"Get Source tensor", - /*retType=*/"::mlir::TypedValue", - /*methodName=*/"getSrc", - /*args=*/(ins)>, - InterfaceMethod< - /*desc=*/"Get mutable source tensor", - /*retType=*/"::mlir::OpOperand&", - /*methodName=*/"getSrcMutable", - /*args=*/(ins)>, - ]; -} - - -#endif // TRITON_OP_INTERFACES diff --git a/third_party/mthreads/include/triton/Dialect/Triton/Transforms/ArithTypeConversion.h b/third_party/mthreads/include/triton/Dialect/Triton/Transforms/ArithTypeConversion.h deleted file mode 100644 index 1e772f330b..0000000000 --- a/third_party/mthreads/include/triton/Dialect/Triton/Transforms/ArithTypeConversion.h +++ /dev/null @@ -1,18 +0,0 @@ -#ifndef TRITON_DIALECT_TRITON_TRANSFORMS_ARITH_TYPE_CONVERSION_H_ -#define TRITON_DIALECT_TRITON_TRANSFORMS_ARITH_TYPE_CONVERSION_H_ -#include "mlir/Transforms/DialectConversion.h" - -namespace mlir::triton { - -/** - * @brief Provides helper patterns for converting arith operations using a type - * converter. - * - * Note at of the time of writing this isn't provided in upstream mlir. - */ -void populateArithTypeConversions(const TypeConverter &converter, - RewritePatternSet &patterns); - -} // namespace mlir::triton - -#endif // TRITON_DIALECT_TRITON_TRANSFORMS_ARITH_TYPE_CONVERSION_H_ diff --git a/third_party/mthreads/include/triton/Dialect/Triton/Transforms/CMakeLists.txt b/third_party/mthreads/include/triton/Dialect/Triton/Transforms/CMakeLists.txt deleted file mode 100644 index 372a9ec11e..0000000000 --- a/third_party/mthreads/include/triton/Dialect/Triton/Transforms/CMakeLists.txt +++ /dev/null @@ -1,3 +0,0 @@ -set(LLVM_TARGET_DEFINITIONS Passes.td) -mlir_tablegen(Passes.h.inc -gen-pass-decls -name Triton) -add_public_tablegen_target(TritonTransformsIncGen) diff --git a/third_party/mthreads/include/triton/Dialect/Triton/Transforms/FunctionTypeConversion.h b/third_party/mthreads/include/triton/Dialect/Triton/Transforms/FunctionTypeConversion.h deleted file mode 100644 index 77940bb417..0000000000 --- a/third_party/mthreads/include/triton/Dialect/Triton/Transforms/FunctionTypeConversion.h +++ /dev/null @@ -1,19 +0,0 @@ -#ifndef TRITON_DIALECT_TRITON_TRANSFORMS_FUNCTION_TYPE_CONVERSION_H_ -#define TRITON_DIALECT_TRITON_TRANSFORMS_FUNCTION_TYPE_CONVERSION_H_ -#include "mlir/Transforms/DialectConversion.h" - -namespace mlir::triton { - -/** - * @brief Provides helper patterns for converting triton function operations - * using a type converter. - * - * Note we cannot use upstream passes for this because they are unaware of - * tt.call and tt.return. - */ -void populateFunctionTypeConversions(const TypeConverter &converter, - RewritePatternSet &patterns); - -} // namespace mlir::triton - -#endif // TRITON_DIALECT_TRITON_TRANSFORMS_FUNCTION_TYPE_CONVERSION_H_ diff --git a/third_party/mthreads/include/triton/Dialect/Triton/Transforms/LoopPeeling.h b/third_party/mthreads/include/triton/Dialect/Triton/Transforms/LoopPeeling.h deleted file mode 100644 index 38efd6b134..0000000000 --- a/third_party/mthreads/include/triton/Dialect/Triton/Transforms/LoopPeeling.h +++ /dev/null @@ -1,18 +0,0 @@ -#ifndef TRITON_DIALECT_TRITON_TRANSFORMS_LOOP_PEELING_H_ -#define TRITON_DIALECT_TRITON_TRANSFORMS_LOOP_PEELING_H_ - -#include "mlir/Dialect/SCF/IR/SCF.h" - -namespace mlir { -namespace triton { - -// Peel the single last iteration of the loop. -void peelLoopEpilogue( - scf::ForOp forOp, - function_ref - processPeeledOp = nullptr); - -} // namespace triton -} // namespace mlir - -#endif // TRITON_DIALECT_TRITON_TRANSFORMS_LOOP_PEELING_H_ diff --git a/third_party/mthreads/include/triton/Dialect/Triton/Transforms/Passes.h b/third_party/mthreads/include/triton/Dialect/Triton/Transforms/Passes.h deleted file mode 100644 index 5d254bf830..0000000000 --- a/third_party/mthreads/include/triton/Dialect/Triton/Transforms/Passes.h +++ /dev/null @@ -1,19 +0,0 @@ -#ifndef TRITON_DIALECT_TRITON_TRANSFORMS_PASSES_H_ -#define TRITON_DIALECT_TRITON_TRANSFORMS_PASSES_H_ - -#include "mlir/Pass/Pass.h" - -namespace mlir { -namespace triton { - -// Generate the pass class declarations. -#define GEN_PASS_DECL -#include "triton/Dialect/Triton/Transforms/Passes.h.inc" - -#define GEN_PASS_REGISTRATION -#include "triton/Dialect/Triton/Transforms/Passes.h.inc" - -} // namespace triton -} // namespace mlir - -#endif diff --git a/third_party/mthreads/include/triton/Dialect/Triton/Transforms/Passes.td b/third_party/mthreads/include/triton/Dialect/Triton/Transforms/Passes.td deleted file mode 100644 index 3744f8ad07..0000000000 --- a/third_party/mthreads/include/triton/Dialect/Triton/Transforms/Passes.td +++ /dev/null @@ -1,93 +0,0 @@ -#ifndef TRITON_PASSES -#define TRITON_PASSES - -include "mlir/Pass/PassBase.td" - -def TritonCombineOps : Pass { - let summary = "combine ops"; - let description = [{ - This pass aims to optimize the five following patterns: - - `dot(a, b, 0) + c => dot(a, b, c)` - - - `addptr(addptr(ptr, idx0), idx1) => addptr(ptr, AddI(idx0, idx1))` - - - `select(cond, load(ptrs, broadcast(cond), ???), other) => - load(ptrs, broadcast(cond), other)` - - - `broadcast(constant) => reshaped_constant` - - `torch.sum(x[:,:,None].expand(-1,-1,n) * y[None,:,:].expand(m,-1,-1),1) - => dot(x,y,splat(0))` - }]; - - let dependentDialects = ["mlir::arith::ArithDialect"]; -} - -def TritonReorderBroadcast : Pass { - let summary = "Moves broadcast and splat after elementwise operations"; - let description = [{ - The purpose of this pass is to transform: - - `elementwise(broadcast(a)) => broadcast(elementwise(a))` - - `elementwise(splat(a), splat(b), ...) => splat(elementwise(a, b, ...))` - In the event of a match, the broadcast (or splat) operation is delayed - and performed after the ElementWise operation. - }]; - - let dependentDialects = ["mlir::triton::TritonDialect"]; -} - -def TritonRewriteTensorPointer : Pass { - let summary = "Rewrite load/stores with tensor pointers into legacy load/stores"; - let description = [{ - This pass rewrites all load/store semantics initiated by a `tt.make_tensor_ptr` and `tt.advance` into legacy - semantics. After this pass, `tt.make_tensor_ptr` and `tt.advance` will disappear, and it generates logics to compute - the pointer/mask/other for each load/store. - }]; - - let dependentDialects = ["mlir::triton::TritonDialect"]; -} - -def TritonRewriteTensorDescriptorToPointer : Pass { - let summary = "Rewrite load/stores of tensor descriptors into pointer load/stores"; - let description = [{ - This pass rewrites all load/store semantics initiated by a `tt.make_tensor_descriptor` into pointer semantics. After - this pass, `tt.make_tensor_descriptor` will disappear, and it generates logics to compute the pointer/mask/other - for each load/store. - }]; - - let dependentDialects = ["mlir::triton::TritonDialect"]; -} - -def TritonLoopUnroll : Pass { - let summary = "Loop unroller"; - let description = [{ - The pass unrolls a scf loop with tt.loop_unroll_factor attribute. The attribute specialises how many iterations - the loop should be unrolled. - }]; - - let dependentDialects = ["mlir::triton::TritonDialect"]; -} - -def TritonLoopInvariantCodeMotion : Pass { - let summary = "MLIR's LICM plus hoist load ops out of loops with masks."; - let description = [{ - This pass uses MLIR's LICM pass as base. Additionally, it hoists load ops - out of loops that consists of pure/read-only ops. For scf.for loops, it - generates a trip-count check. For scf.while loops, it clones the condition - from the before body. - }]; - - let dependentDialects = ["mlir::triton::TritonDialect"]; -} - -def TritonLoopAwareCSE : Pass<"triton-loop-aware-cse", "mlir::ModuleOp"> { - let summary = "CSE within loop bodies"; - - let description = [{ - The `triton-loop-aware-cse` pass performs recursive common subexpression - elimination within loop bodies. Unlike regular CSE, which is a single-pass - greedy algorithm, this pass can recursively eliminate loop iteration - arguments and subcomputations that always have the same value. - }]; -} - -#endif diff --git a/third_party/mthreads/include/triton/Dialect/TritonGPU/CMakeLists.txt b/third_party/mthreads/include/triton/Dialect/TritonGPU/CMakeLists.txt deleted file mode 100644 index 9f57627c32..0000000000 --- a/third_party/mthreads/include/triton/Dialect/TritonGPU/CMakeLists.txt +++ /dev/null @@ -1,2 +0,0 @@ -add_subdirectory(IR) -add_subdirectory(Transforms) diff --git a/third_party/mthreads/include/triton/Dialect/TritonGPU/IR/CMakeLists.txt b/third_party/mthreads/include/triton/Dialect/TritonGPU/IR/CMakeLists.txt deleted file mode 100644 index 5ce12376ba..0000000000 --- a/third_party/mthreads/include/triton/Dialect/TritonGPU/IR/CMakeLists.txt +++ /dev/null @@ -1,45 +0,0 @@ -set(MLIR_BINARY_DIR ${CMAKE_BINARY_DIR}) - -set(LLVM_TARGET_DEFINITIONS TritonGPUOps.td) -if(FLAGTREE_MTHREADS_TLE) - set(_TLE_TABLEGEN_DEFS -D__TLE__) -else() - set(_TLE_TABLEGEN_DEFS "") -endif() -mlir_tablegen(Dialect.h.inc -gen-dialect-decls -dialect=ttg) -mlir_tablegen(Dialect.cpp.inc -gen-dialect-defs -dialect=ttg) -mlir_tablegen(Ops.h.inc -gen-op-decls ${_TLE_TABLEGEN_DEFS}) -mlir_tablegen(Ops.cpp.inc -gen-op-defs ${_TLE_TABLEGEN_DEFS}) -mlir_tablegen(Types.h.inc -gen-typedef-decls -typedefs-dialect=ttg) -mlir_tablegen(Types.cpp.inc -gen-typedef-defs -typedefs-dialect=ttg) -add_mlir_doc(TritonGPUDialect TritonGPUDialect dialects/ -gen-dialect-doc) -add_mlir_doc(TritonGPUOps TritonGPUOps dialects/ -gen-op-doc) -add_public_tablegen_target(TritonGPUTableGen) - -set(LLVM_TARGET_DEFINITIONS TritonGPUAttrDefs.td) -mlir_tablegen(AttrInterfaces.h.inc -gen-attr-interface-decls) -mlir_tablegen(AttrInterfaces.cpp.inc -gen-attr-interface-defs) -mlir_tablegen(AttrDefs.h.inc -gen-attrdef-decls) - -set(LLVM_TARGET_DEFINITIONS TritonGPUAttrImpls.td) -mlir_tablegen(AttrDefs.cpp.inc -gen-attrdef-defs) -add_public_tablegen_target(TritonGPUAttrDefsIncGen) - -set(LLVM_TARGET_DEFINITIONS TritonGPUEnums.td) -mlir_tablegen(OpsEnums.h.inc -gen-enum-decls) -mlir_tablegen(OpsEnums.cpp.inc -gen-enum-defs) -add_public_tablegen_target(TritonGPUOpsEnumsIncGen) - -set(LLVM_TARGET_DEFINITIONS CGAEncodingAttr.td) -mlir_tablegen(CGAEncodingAttr.h.inc -gen-attrdef-decls) -add_public_tablegen_target(TritonGPUCGAAttrIncGen) - -set(LLVM_TARGET_DEFINITIONS TritonGPUTypeInterfaces.td) -mlir_tablegen(TypeInterfaces.h.inc -gen-type-interface-decls) -mlir_tablegen(TypeInterfaces.cpp.inc -gen-type-interface-defs) -add_public_tablegen_target(TritonGPUTypeInterfacesIncGen) - -set(LLVM_TARGET_DEFINITIONS TritonGPUOpInterfaces.td) -mlir_tablegen(OpInterfaces.h.inc -gen-op-interface-decls) -mlir_tablegen(OpInterfaces.cpp.inc -gen-op-interface-defs) -add_public_tablegen_target(TritonGPUOpInterfacesIncGen) diff --git a/third_party/mthreads/include/triton/Dialect/TritonGPU/IR/TritonGPUDialect.td b/third_party/mthreads/include/triton/Dialect/TritonGPU/IR/TritonGPUDialect.td deleted file mode 100644 index 3169dc451f..0000000000 --- a/third_party/mthreads/include/triton/Dialect/TritonGPU/IR/TritonGPUDialect.td +++ /dev/null @@ -1,41 +0,0 @@ -#ifndef TRITONGPU_DIALECT -#define TRITONGPU_DIALECT - -include "mlir/IR/OpBase.td" - -def TritonGPU_Dialect : Dialect { - let name = "ttg"; - - let cppNamespace = "::mlir::triton::gpu"; - - let hasOperationAttrVerify = 1; - - let description = [{ - Triton GPU Dialect. - }]; - - let dependentDialects = [ - "triton::TritonDialect", - "mlir::gpu::GPUDialect", - ]; - - let extraClassDeclaration = [{ - void registerTypes(); - - LinearLayout toLinearLayout(ArrayRef shape, Attribute layout); - LinearEncodingAttr toLinearEncoding(ArrayRef shape, Attribute layout); - - static int getNumCTAs(ModuleOp mod); - static int getThreadsPerWarp(ModuleOp mod); - - private: - LinearLayoutCache llCache; - LinearEncodingCache leCache; - }]; - - let useDefaultTypePrinterParser = 1; - let useDefaultAttributePrinterParser = 1; - let usePropertiesForAttributes = 1; -} - -#endif diff --git a/third_party/mthreads/include/triton/Dialect/TritonGPU/IR/TritonGPUOpInterfaces.td b/third_party/mthreads/include/triton/Dialect/TritonGPU/IR/TritonGPUOpInterfaces.td deleted file mode 100644 index 3862b7f474..0000000000 --- a/third_party/mthreads/include/triton/Dialect/TritonGPU/IR/TritonGPUOpInterfaces.td +++ /dev/null @@ -1,29 +0,0 @@ -#ifndef TRITONGPU_OP_INTERFACES -#define TRITONGPU_OP_INTERFACES - -include "mlir/IR/OpBase.td" - -def UpcastFpOpInterface : OpInterface<"UpcastFpOpInterface"> { - let description = [{ - This interface is for operations that upcast floating-point numbers. - }]; - - let cppNamespace = "::mlir::triton::gpu"; - - let methods = [ - InterfaceMethod< - /*desc=*/"Infer destination encoding", - /*retType=*/"mlir::Attribute", - /*methodName=*/"inferDstEncoding", - /*args=*/(ins "unsigned":$opIdx, "mlir::Attribute":$srcEnc) - >, - InterfaceMethod< - /*desc=*/"Infer operand encoding from dst encoding", - /*retType=*/"mlir::Attribute", - /*methodName=*/"inferSrcEncoding", - /*args=*/(ins "unsigned":$opIdx, "mlir::Attribute":$dstEnc) - > - ]; -} - -#endif // TRITONGPU_OP_INTERFACES diff --git a/third_party/mthreads/include/triton/Dialect/TritonGPU/IR/TritonGPUTypeInterfaces.td b/third_party/mthreads/include/triton/Dialect/TritonGPU/IR/TritonGPUTypeInterfaces.td deleted file mode 100644 index a0415b62c6..0000000000 --- a/third_party/mthreads/include/triton/Dialect/TritonGPU/IR/TritonGPUTypeInterfaces.td +++ /dev/null @@ -1,23 +0,0 @@ -#ifndef TRITON_GPU_TYPE_INTERFACES -#define TRITON_GPU_TYPE_INTERFACES - -include "mlir/IR/OpBase.td" - -// Interface dynamically attached to RankedTensorType and MemDescType. -def TTG_TensorOrMemDesc : TypeInterface<"TensorOrMemDesc"> { - let cppNamespace = "::mlir::triton::gpu"; - let methods = [ - InterfaceMethod<"Returns the encoding of the tensor or memory descriptor", - "mlir::Attribute", "getEncoding", (ins)>, - InterfaceMethod<"Returns element type", - "mlir::Type", "getElementType", (ins)>, - InterfaceMethod<"Returns the type shape", - "llvm::ArrayRef", "getShape", (ins)>, - InterfaceMethod<"Returns the tensor or buffer rank", - "int64_t", "getRank", (ins)>, - InterfaceMethod<"Returns the element type bit width", - "int64_t", "getElementTypeBitWidth", (ins)>, - ]; -} - -#endif // TRITON_GPU_TYPE_INTERFACES diff --git a/third_party/mthreads/include/triton/Dialect/TritonGPU/IR/Types.h b/third_party/mthreads/include/triton/Dialect/TritonGPU/IR/Types.h deleted file mode 100644 index cfad8be199..0000000000 --- a/third_party/mthreads/include/triton/Dialect/TritonGPU/IR/Types.h +++ /dev/null @@ -1,14 +0,0 @@ -#ifndef TRITONGPU_IR_TYPES_H_ -#define TRITONGPU_IR_TYPES_H_ - -#include "mlir/IR/BuiltinTypes.h" -#include "mlir/IR/TypeSupport.h" -#include "mlir/IR/Types.h" -#include "triton/Dialect/TritonGPU/IR/Attributes.h" - -#define GET_TYPEDEF_CLASSES -#include "triton/Dialect/TritonGPU/IR/Types.h.inc" - -#include "triton/Dialect/TritonGPU/IR/TypeInterfaces.h.inc" - -#endif // TRITON_IR_TYPES_H_ diff --git a/third_party/mthreads/include/triton/Dialect/TritonGPU/Transforms/CMakeLists.txt b/third_party/mthreads/include/triton/Dialect/TritonGPU/Transforms/CMakeLists.txt deleted file mode 100644 index 6be94d1a8a..0000000000 --- a/third_party/mthreads/include/triton/Dialect/TritonGPU/Transforms/CMakeLists.txt +++ /dev/null @@ -1,3 +0,0 @@ -set(LLVM_TARGET_DEFINITIONS Passes.td) -mlir_tablegen(Passes.h.inc -gen-pass-decls -name TritonGPU) -add_public_tablegen_target(TritonGPUTransformsIncGen) diff --git a/third_party/mthreads/include/triton/Dialect/TritonGPU/Transforms/DecomposeScaledBlocked.h b/third_party/mthreads/include/triton/Dialect/TritonGPU/Transforms/DecomposeScaledBlocked.h deleted file mode 100644 index f06f85e58a..0000000000 --- a/third_party/mthreads/include/triton/Dialect/TritonGPU/Transforms/DecomposeScaledBlocked.h +++ /dev/null @@ -1,47 +0,0 @@ -#include "mlir/IR/PatternMatch.h" -#include "triton/Dialect/Triton/IR/Dialect.h" - -namespace mlir::triton::gpu { - -class DecomposeScaledBlocked : public OpRewritePattern { -public: - DecomposeScaledBlocked(MLIRContext *context, PatternBenefit benefit) - : OpRewritePattern(context, benefit) {} - - LogicalResult matchAndRewrite(DotScaledOp scaledDotOp, - PatternRewriter &rewriter) const override; - -protected: - FloatType getComputeType(ScaleDotElemType aType, ScaleDotElemType bType, - PatternRewriter &rewriter) const; - TypedValue scaleTo16(PatternRewriter &rewriter, - TypedValue scale, - FloatType computeType) const; - TypedValue - broadcastScale(PatternRewriter &rewriter, DotScaledOp scaledDotOp, - ModuleOp mod, TypedValue scale, - int dim) const; - TypedValue maskNan(PatternRewriter &rewriter, - DotScaledOp scaledDotOp, - TypedValue mxfp, - TypedValue scale, - int dim) const; - virtual TypedValue scaleArg(PatternRewriter &rewriter, - DotScaledOp scaledDotOp, - int opIdx, - FloatType computeType) const; - TypedValue - cvtDotOperand(PatternRewriter &rewriter, DotScaledOp scaledDotOp, int opIdx, - TypedValue v) const; - TypedValue - extendAndBroadcastScale(PatternRewriter &rewriter, DotScaledOp scaledDotOp, - TypedValue &scale, - FloatType computeType, RankedTensorType dstType, - int opIdx) const; - static SmallVector getTransposeOrder(int rank); -}; - -void populateDecomposeScaledBlockedPatterns(mlir::RewritePatternSet &patterns, - int benefit); - -} // namespace mlir::triton::gpu diff --git a/third_party/mthreads/include/triton/Dialect/TritonGPU/Transforms/LayoutPropagationUtility.h b/third_party/mthreads/include/triton/Dialect/TritonGPU/Transforms/LayoutPropagationUtility.h deleted file mode 100644 index b289de5593..0000000000 --- a/third_party/mthreads/include/triton/Dialect/TritonGPU/Transforms/LayoutPropagationUtility.h +++ /dev/null @@ -1,21 +0,0 @@ -#ifndef TRITON_DIALECT_TRITONGPU_TRANSFORMS_LAYOUT_PROPAGATION_UTILITY_H_ -#define TRITON_DIALECT_TRITONGPU_TRANSFORMS_LAYOUT_PROPAGATION_UTILITY_H_ - -#include "triton/Dialect/Triton/IR/Dialect.h" -#include "triton/Dialect/TritonGPU/IR/Attributes.h" -#include "triton/Tools/LinearLayout.h" -#include - -namespace mlir::triton::gpu { - -// Given the result |dstLayout|, infer the source layout that we should use for -// global load if we propagate through op def chain of |defOp|. Returns -// std::nullopt if fails to infer or cannot reach a global load. -std::optional> -inferSourceLoadLayout(const LinearLayout &dstLayout, Operation *defOp); -std::optional> -inferSourceLoadLayout(LinearEncodingAttr dstLayout, Operation *defOp); - -} // namespace mlir::triton::gpu - -#endif // TRITON_DIALECT_TRITONGPU_TRANSFORMS_LAYOUT_PROPAGATION_UTILITY_H_ diff --git a/third_party/mthreads/include/triton/Dialect/TritonGPU/Transforms/MMAv5PipelineUtility.h b/third_party/mthreads/include/triton/Dialect/TritonGPU/Transforms/MMAv5PipelineUtility.h deleted file mode 100644 index 58e5290c29..0000000000 --- a/third_party/mthreads/include/triton/Dialect/TritonGPU/Transforms/MMAv5PipelineUtility.h +++ /dev/null @@ -1,83 +0,0 @@ -#ifndef TRITON_TRITONGPU_TRANSFORMS_MMAV5PIPELINEUTILITY_H_ -#define TRITON_TRITONGPU_TRANSFORMS_MMAV5PIPELINEUTILITY_H_ - -#include "triton/Dialect/TritonNvidiaGPU/IR/Dialect.h" - -namespace mlir { - -class OpBuilder; -class DominanceInfo; - -namespace scf { -class ForOp; -} // namespace scf -namespace triton::nvidia_gpu { - -//===----------------------------------------------------------------------===// -// MMA Pipeline Analysis -//===----------------------------------------------------------------------===// - -// Given an MMAv5 operation in a loop, determine if its accumulator can be -// multibuffered. -bool isAccMultibufferingPossible(MMAv5OpInterface mma, scf::ForOp forOp); - -// Returns true if the MMA operation requires acc multi-buffering when -// pipelined. -bool requiresAccMultiBuffering(MMAv5OpInterface mma, scf::ForOp forOp); - -// Returns true if there are loads from tmem after the MMA operation. -bool hasLoadsAfterMMA(MMAv5OpInterface mma, scf::ForOp forOp); - -// Helper class to determine if the operands of an MMA operation are -// pipelineable. -class MMAv5PipelineableOperandsHelper { -public: - MMAv5PipelineableOperandsHelper( - MMAv5OpInterface mmaOp, scf::ForOp forOp, - std::function isLoadToBePipelined) - : mmaOp(mmaOp), forOp(forOp), isLoadToBePipelined(isLoadToBePipelined) { - run(); - } - - bool isPipelineable = false; - // If true, the existing operand loads are all been found and their - // pipelineability has been determined. - bool isOperandsStateDetermined = false; - SmallVector unpipelineableOperandDefs; - -private: - MMAv5OpInterface mmaOp; - scf::ForOp forOp; - std::function isLoadToBePipelined; - void run(); - bool isOperandPipelineable(Value v, Operation *&foundDef); -}; - -bool areScalesPipelineable(TCGen5MMAScaledOp scaledOp, scf::ForOp forOp); -bool isOperandPipelineableBase( - Value v, scf::ForOp forOp, Operation *&foundDef, - std::function isPipelineable = - [](Operation *) { return false; }, - std::function isLoadToBePipelined = - [](Operation *) { return false; }); - -//===----------------------------------------------------------------------===// -// MMA Pipeline Rewriters -//===----------------------------------------------------------------------===// - -// Create a new TMEMAllocOp to use for the pipelined MMA operation. It is -// optionally multi-buffered based on the number of stages. -TMEMAllocOp createTMemAlloc(OpBuilder &builder, TMEMAllocOp oldTMemAllocOp, - bool multiBufferred, int numStages); - -// Return true if the accumulator of an mma in subsequent iterations is either -// independent from the previous iteration (overwritten) or completely reused, -// without read-modify-write. -// Otherwise, we can not pipeline the MMA, as we need to insert a wait after the -// mma to read back the accumulator for RMW. -bool hasAccReadModifyWrite(MMAv5OpInterface mma, scf::ForOp forOp); - -} // namespace triton::nvidia_gpu -} // namespace mlir - -#endif // TRITON_TRITONGPU_TRANSFORMS_MMAV5PIPELINEUTILITY_H_ diff --git a/third_party/mthreads/include/triton/Dialect/TritonGPU/Transforms/Partition.h b/third_party/mthreads/include/triton/Dialect/TritonGPU/Transforms/Partition.h deleted file mode 100644 index 6c5b287f0c..0000000000 --- a/third_party/mthreads/include/triton/Dialect/TritonGPU/Transforms/Partition.h +++ /dev/null @@ -1,127 +0,0 @@ -#ifndef TRITON_TRITONGPU_TRANSFORM_PIPELINE_PARTITION_H_ -#define TRITON_TRITONGPU_TRANSFORM_PIPELINE_PARTITION_H_ - -#include "mlir/Support/LLVM.h" -#include "llvm/ADT/ArrayRef.h" -#include "llvm/ADT/DenseMap.h" -#include "llvm/ADT/SmallVector.h" - -namespace mlir { -class Operation; -class OpOperand; -class OpResult; -class Region; -namespace scf { -class ForOp; -} // namespace scf -} // namespace mlir - -//===----------------------------------------------------------------------===// -// PartitionSet -//===----------------------------------------------------------------------===// - -namespace mlir::triton::gpu { -// A partition has a stage and contains some operation. The stage of a -// partition determines how many cycles the partition's outputs are buffered -// relative to its consumers. -class Partition { -public: - Partition(int idx, int stage) : idx(idx), stage(stage) { - assert(idx >= 0 && "A partition index must be nonnegative."); - } - - int getIndex() const { return idx; } - int getStage() const { return stage; } - ArrayRef getOps() const { return ops; } - void addOp(Operation *op) { ops.push_back(op); } - bool hasOp(Operation *op) const; - bool empty() const { return ops.empty(); } - - // Iterate the inputs of the partition. Input values are those that originate - // from a different partition or a previous iteration of the current - // partition. E.g. partition B(i) may have inputs from A(i) or B(i-1). Note - // that the same value may be visited more than once. - void iterateInputs(scf::ForOp loop, - function_ref callback) const; - // Iterate the outputs of the partition. Output values are those that are - // consumed by a different partition or a future iteration of the current - // partition. E.g. partition A(i) may have outputs to B(i) or A(i+1). Note - // that the same value may be visited more than once. - void - iterateOutputs(scf::ForOp loop, - function_ref callback) const; - // Iterate the defining ops of the inputs to the partition in the current and - // previous iterations, including the distance in the past. - void iterateDefs(scf::ForOp loop, - function_ref callback) const; - // Iterate the uses of all outputs of the partition in the current iteration - // and in future iterations, including the distance in the future. - void iterateUses( - scf::ForOp loop, - function_ref callback) const; - -private: - void setIndex(int idx) { this->idx = idx; } - - // The partition number. - int idx; - // The stage of the partition. - int stage; - // The ops in the partition. - SmallVector ops; -}; - -// A partition set divides a loop into multiple partitions. Ops in a loop are -// assigned at most one partition. A partition set represents asynchronous -// execution of the loop body, where partitions may execute simultaneously. -class PartitionSet { -public: - // Get WarpSpecialization tag - int getTag() const { return tag; } - - // Create a new partition with a stage. - Partition *addPartition(unsigned stage); - - // Get the partition at the index. - Partition *getPartition(unsigned idx); - // Get the partition at the index. - const Partition *getPartition(unsigned idx) const; - // Return an iterator range over the partitions. - auto getPartitions() { return llvm::make_pointee_range(partitions); } - // Return an iterator range over the partitions. - auto getPartitions() const { return llvm::make_pointee_range(partitions); } - // Get the number of partitions. - unsigned getNumPartitions() const { return partitions.size(); } - - // Deserialize a partition set from an `scf.for` op using the attributes - // tagged on operations in its body. - static FailureOr fromLoop(scf::ForOp loop); - - // Debug dump the partition set. - LLVM_DUMP_METHOD void dump() const; - - // Utility to be used when the op is known to belong to one partition - Partition *getPartition(Operation *op); - -private: - // WarpSpecialization tag - int tag; - // Partitions are numbered [0, N). - SmallVector> partitions; -}; - -// Annotate the op with the partition index or indices, and add the op -// to the partitions it belongs to. -void setPartition(Operation *op, Partition *partition); -void setPartition(Operation *op, const SetVector &partitions); -// Annotate the op with the partition indices. It should only be used in a pass -// which does not work with Partition instances and iterate* functions, since -// it does not keep the op attributes and the op list of a partition in sync. -void setPartition(Operation *op, const SetVector &partitionIds); -void setPartitionOutputs(Operation *op, - ArrayRef> partitionOutputsIds); -void setWarpSpecializeTag(Operation *op, int tag); - -} // namespace mlir::triton::gpu - -#endif // TRITON_TRITONGPU_TRANSFORM_PIPELINE_PARTITION_H_ diff --git a/third_party/mthreads/include/triton/Dialect/TritonGPU/Transforms/PartitionBuilder.h b/third_party/mthreads/include/triton/Dialect/TritonGPU/Transforms/PartitionBuilder.h deleted file mode 100644 index baa16421c1..0000000000 --- a/third_party/mthreads/include/triton/Dialect/TritonGPU/Transforms/PartitionBuilder.h +++ /dev/null @@ -1,49 +0,0 @@ -#ifndef TRITON_TRITONGPU_TRANSFORMS_PARTITIONBUILDER_H -#define TRITON_TRITONGPU_TRANSFORMS_PARTITIONBUILDER_H - -#include "mlir/IR/ImplicitLocOpBuilder.h" -#include "llvm/ADT/SetVector.h" - -namespace mlir::triton::gpu { - -class Partition; - -using StageCluster = std::optional>; - -// Get the stage and cluster for an operation, if it has one assigned. -void setStageCluster(OpBuilder &b, Operation *op, StageCluster stageCluster); -StageCluster getStageCluster(Operation *op); - -struct PartitionBuilder : public ImplicitLocOpBuilder { - using ImplicitLocOpBuilder::ImplicitLocOpBuilder; - - Value intCst(int value, unsigned width = 32); - Value boolCst(bool value); - - void assignPartition(Operation *op, Partition &partition); - - template - auto createInto(Partition &partition, StageCluster stageCluster, - Args &&...args) { - auto op = create(std::forward(args)...); - assignPartition(op, partition); - setStageCluster(*this, op, stageCluster); - return op; - } -}; - -template -OpT createInto(OpBuilder &b, Location loc, - std::optional> partitionSet, - StageCluster stageCluster, Args &&...args) { - auto op = OpT::create(b, loc, std::forward(args)...); - if (partitionSet) { - setPartition(op, *partitionSet); - setStageCluster(b, op, stageCluster); - } - return op; -} - -} // namespace mlir::triton::gpu - -#endif // TRITON_TRITONGPU_TRANSFORMS_PARTITIONBUILDER_H diff --git a/third_party/mthreads/include/triton/Dialect/TritonGPU/Transforms/PipelineExpander.h b/third_party/mthreads/include/triton/Dialect/TritonGPU/Transforms/PipelineExpander.h deleted file mode 100644 index 4851bfe001..0000000000 --- a/third_party/mthreads/include/triton/Dialect/TritonGPU/Transforms/PipelineExpander.h +++ /dev/null @@ -1,111 +0,0 @@ -// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. -// See https://llvm.org/LICENSE.txt for license information. -// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception -#ifndef TRITON_DIALECT_TRITONGPU_TRANSFORMS_PIPELINE_H_ -#define TRITON_DIALECT_TRITONGPU_TRANSFORMS_PIPELINE_H_ - -// This is a fork of upstream pipeline transformation. This will be merged back -// upstream once we have a stable solution. - -#include "mlir/Support/LLVM.h" -#include "mlir/Support/LogicalResult.h" -#include "llvm/ADT/ArrayRef.h" - -namespace mlir { - -class RewriterBase; -class Operation; -class Value; - -namespace scf { -class ForOp; -} - -namespace triton { - -/// Options to dictate how loops should be pipelined. -struct PipeliningOption { - /// Lambda returning all the operations in the forOp, with their stage, in the - /// order picked for the pipelined loop. - using GetScheduleFnType = std::function> &)>; - GetScheduleFnType getScheduleFn = nullptr; - enum class PipelinerPart { - Prologue, - Kernel, - Epilogue, - }; - /// Lambda called by the pipeliner to allow the user to annotate the IR while - /// it is generated. - /// The callback passes the operation created along with the part of the - /// pipeline and the iteration index. The iteration index is always 0 for the - /// kernel. For the prologue and epilogue, it corresponds to the iteration - /// peeled out of the loop in the range [0, maxStage[. - using AnnotationlFnType = - std::function; - AnnotationlFnType annotateFn = nullptr; - - /// Control whether the epilogue should be peeled out of the loop or - /// operations should be predicated to skip the early stages in the last loop - /// iterations. If the epilogue is predicated; the user needs to provide a - /// lambda to generate the predicated version of operations. - bool peelEpilogue = true; - - /// Control whether the transformation checks that the number of iterations is - /// greater or equal to the number of stages and skip the transformation if - /// this is not the case. If the loop is dynamic and this is set to true the - /// pipeliner will have to predicate operations in the prologue/epilogue. - bool supportDynamicLoops = false; - - /// If set, use this function to emit the predicate stage ops instead of the - /// default one. - using EmitPredicateStageFnType = std::function; - EmitPredicateStageFnType emitPredicateStageFn = nullptr; - - // Callback to predicate operations when the prologue or epilogue are not - // peeled. This takes the original operation, an i1 predicate value and the - // pattern rewriter. It is expected to replace the given operation with - // the predicated equivalent and return it, or return nullptr if the - // predication is impossible. In the latter case, pipelining will fail and - // may leave IR in a partially transformed state. - using PredicateOpFnType = - std::function; - PredicateOpFnType predicateFn = nullptr; - - // TODO: add option to decide if the prologue should be peeled. -}; - -/// Generate a pipelined version of the scf.for loop based on the schedule given -/// as option. This applies the mechanical transformation of changing the loop -/// and generating the prologue/epilogue for the pipelining and doesn't make any -/// decision regarding the schedule. -/// Based on the options the loop is split into several stages. -/// The transformation assumes that the scheduling given by user is valid. -/// For example if we break a loop into 3 stages named S0, S1, S2 we would -/// generate the following code with the number in parenthesis as the iteration -/// index: -/// -/// S0(0) // Prologue -/// S0(1) S1(0) // Prologue -/// scf.for %I = %C0 to %N - 2 { -/// S0(I+2) S1(I+1) S2(I) // Pipelined kernel -/// } -/// S1(N) S2(N-1) // Epilogue -/// S2(N) // Epilogue -/// -/// If `modifiedIR` is provided, it will be set to a value that indicates -/// whether pipelining modified the IR before failing, signaling to the caller -/// whether they can proceed with different transformations. -FailureOr pipelineForLoop(RewriterBase &rewriter, scf::ForOp forOp, - const PipeliningOption &options, - bool *modifiedIR = nullptr); - -Value emitPredicateForStage(RewriterBase &rewriter, Value inductionVar, - Value upperBound, Value step, uint64_t maxStage, - uint64_t stage); - -} // namespace triton -} // namespace mlir - -#endif // TRITON_DIALECT_TRITONGPU_TRANSFORMS_PIPELINE_H_ diff --git a/third_party/mthreads/include/triton/Dialect/TritonGPU/Transforms/PipeliningUtility.h b/third_party/mthreads/include/triton/Dialect/TritonGPU/Transforms/PipeliningUtility.h deleted file mode 100644 index 5700a366fc..0000000000 --- a/third_party/mthreads/include/triton/Dialect/TritonGPU/Transforms/PipeliningUtility.h +++ /dev/null @@ -1,189 +0,0 @@ -#ifndef TRITON_TRITONGPU_TRANSFORMS_PIPELINER_PIPELINING_UTILITY_H_ -#define TRITON_TRITONGPU_TRANSFORMS_PIPELINER_PIPELINING_UTILITY_H_ - -#include "mlir/Dialect/SCF/IR/SCF.h" -#include "triton/Dialect/TritonGPU/IR/Dialect.h" -#include -#include -#include - -namespace mlir { -class DominanceInfo; -class ImplicitLocOpBuilder; -namespace triton { - -static const char *kNumStagesAttrName = "tt.num_stages"; -static const char *kDisallowAccMultiBufferAttrName = - "tt.disallow_acc_multi_buffer"; -static const char *kWarpSpecializeAttrName = "tt.warp_specialize"; -static const char *kLoopStageAttrName = "loop.stage"; -static const char *kLoopClusterAttrName = "loop.cluster"; -static const char *kScheduledMaxStageAttrName = "tt.scheduled_max_stage"; -class CoarseSchedule; -class ModuleAxisInfoAnalysis; -//===----------------------------------------------------------------------===// -// Hoisting Utilities -//===----------------------------------------------------------------------===// - -// By default, an operation can be hoisted if it is pure scalar operation. -bool isPureScalarOp(Operation *op); - -// Given a set of values and a reference operation, return true if all of the -// values dominate the reference operation OR a set of "trivial" operations can -// be moved before the reference operation such that the value set dominates the -// reference operation. -// -// Returns false if it is not possible to make the values dominate the reference -// operation. The function determines "trivial"-ness with the given callback. -// By default, it determines that memory-effect-free and scalar operations are -// trivial. -bool getDominatingValueSetOpsToHoist( - DominanceInfo &domInfo, Operation *refOp, ArrayRef valueSet, - llvm::SetVector &toHoist, - function_ref canHoist = isPureScalarOp, - function_ref canUseArg = [](BlockArgument) { - return false; - }); - -// Hoist the given set of operations above the reference operation. -void hoistOpsBefore(Operation *refOp, - const llvm::SetVector &toHoist); -// Hoist the given set of operations before the iterator. -void hoistOpsBefore(Block *block, Block::iterator it, - const llvm::SetVector &toHoist); - -//===----------------------------------------------------------------------===// -// Sinking Utilities -//===----------------------------------------------------------------------===// - -// Sink a value redefinition into a block, provided that the block is dominated -// by `in` and postdominated by `out`. -Value sinkValueRedefinition(RewriterBase &rewriter, Value in, Value out, - Block *block); - -//===----------------------------------------------------------------------===// -// Loop Pipelining Utilities -//===----------------------------------------------------------------------===// - -bool loopHasDistGreaterThanOne(scf::ForOp forOp); -bool isOuterLoop(scf::ForOp forOp); - -/// Function to mask operations during scheduling. -Operation *predicateOp(RewriterBase &rewriter, Operation *op, Value pred); - -/// Wrap the operation into a MaskOp using the provided predicate, enabling high -/// level predication abstraction during pipelining. -Operation *wrapInMaskOp(RewriterBase &rewriter, Operation *op, Value pred); - -// Utilize high level predication abstraction to perform optimizations before -// lowering to predicated operations -void resolveMaskOp(ModuleOp moduleOp); - -// Return true if the given ForOp has the attribute -// `tt.disallow_acc_multi_buffer` set to true. -bool getDisallowAccMultiBuffer(scf::ForOp forOp); - -// Return the definition of the given value. If the value is a loop-carried -// dependency, return the definition and the distance to it. -std::pair getDefinitionAndDistance(scf::ForOp forOp, - Value value); -// Return the defining op of the given value, if the Value is an argument of the -// loop return the associated defining op in the loop and its distance to the -// Value. -std::pair getDefiningOpAndDistance(scf::ForOp forOp, - Value value); - -// Return maximum length of the vectorized copy between registers and shared -// memory for the given tensor type and shared encoding. -int getCopyVecBytes(RankedTensorType registerTy, - gpu::SharedEncodingTrait sharedEnc); - -bool canBeConvertedToAsyncLoad( - triton::LoadOp loadOp, triton::ModuleAxisInfoAnalysis &axisInfoAnalysis); - -// Serialize the latencies of the operations in the loops into the latency -// attribute. -void serializeLatencies(ModuleOp module, DenseMap &opLatency); - -// Serialize the self latencies of the operations in the loops into the -// self_latency attribute. -void serializeSelfLatencies(ModuleOp module, - DenseMap &opSelfLatency); - -// Deserialize the latencies of the operations in the loops from the attribute. -DenseMap deserializeLatencies(Operation *op); - -// Create an allocation for multibuffered scalars. -Value createScalarAlloc(ImplicitLocOpBuilder &rewriter, Type type, - unsigned numBuffers); -// Create an allocation and init the mbarriers. -Value createBarrierAlloc(Operation *op, int numBarriers, int arriveCount = 1); -// Create an allocation that can hold distance number of tensor shapes. -Value createAlloc(Operation *insertBefore, RankedTensorType ty, Location loc, - gpu::SharedEncodingTrait sharedEnc, unsigned distance); - -// Determine if the operation is a TMA load. -bool isTMALoad(Operation *op); - -// Determine if the operation can be lowered to an async load. -bool canBeAsyncLoad(Operation *op); - -// Look for consecutive wait ops and combine them into a single wait op. -void combineRedundantWaitOps( - llvm::SmallSetVector &waitOps); - -// Get the type of the view of a multi-buffered tensor value. -gpu::MemDescType getBufferViewType(gpu::MemDescType allocTy, - bool mutableMemory = true); - -// Get a mutable, multi-buffered version of the given memdesc type, with -// multiplicity "depth". -gpu::MemDescType getMultiBufferedType(gpu::MemDescType memDescType, - int32_t depth); - -// Get a generic shared encoding for a tensor. -gpu::SharedEncodingTrait getSharedEncoding(RankedTensorType ty); -// Get a shared encoding for a tensor based on its uses. -gpu::SharedEncodingTrait getSharedEncoding(Operation *loadOp); - -// Get the number of stages to pipeline the loop with, if it is explicitly -// specified. -int getNumStagesOrDefault(scf::ForOp forOp, int defaultNumStages); - -// Given a result of MemDescIndex, or Alloca, create a MemDescIndex with a -// single buffer slice (leading dimension equal to 1), at the given index. -TypedValue -createSingleBufferView(OpBuilder &builder, Value alloc, Value idx); -// Given a result of MemDescIndex, or Alloca, create a MemDescIndex with a -// single buffer slice (leading dimension equal to 1), at the given index. -TypedValue -createSingleBufferView(OpBuilder &builder, Value alloc, int idx); - -Value createIncrementModulo(OpBuilder &builder, Location loc, Value counter, - Value modulus, Value zero, Value one, - Value *outWrapCond = nullptr); - -scf::ForOp lowerTMADescriptors(scf::ForOp forOp, CoarseSchedule &schedule); - -DenseSet -getTopLevelUsersInLoop(Operation *op, scf::ForOp forOp, - std::function filter = nullptr); - -// Return the "first" op in terms of the stage and cluser ordering -Operation * -getFirstUseOfPipelinedOp(ArrayRef ops, scf::ForOp forOp, - CoarseSchedule &schedule, - std::function filterUse = nullptr); - -// Return the "last" op in terms of the stage and cluser ordering -Operation * -getLastUseOfPipelinedOp(ArrayRef ops, scf::ForOp forOp, - CoarseSchedule &schedule, - std::function filterUse = nullptr); - -// Clean up attributes passing over schedules across stages in pipelining -void removePipeliningAttributes(ModuleOp moduleOp); -} // namespace triton -} // namespace mlir - -#endif // TRITON_TRITONGPU_TRANSFORMS_PIPELINER_PIPELINING_UTILITY_H_ diff --git a/third_party/mthreads/include/triton/Dialect/TritonGPU/Transforms/WarpSpecialization.h b/third_party/mthreads/include/triton/Dialect/TritonGPU/Transforms/WarpSpecialization.h deleted file mode 100644 index afb7dde2c1..0000000000 --- a/third_party/mthreads/include/triton/Dialect/TritonGPU/Transforms/WarpSpecialization.h +++ /dev/null @@ -1,24 +0,0 @@ -#ifndef TRITON_TRITONGPU_TRANSFORM_PIPELINE_WARPSPECIALIZATION_H_ -#define TRITON_TRITONGPU_TRANSFORM_PIPELINE_WARPSPECIALIZATION_H_ - -#include "mlir/Support/LogicalResult.h" - -namespace mlir { -namespace scf { -class ForOp; -} // namespace scf -namespace triton::gpu { -// This is the final step to prepare a loop for warp specialization. This takes -// a loop with a partition schedule and rewrites the loop such that all SSA -// dependencies between partitions are passed through shared memory and -// multibuffers them according to partition stages. -LogicalResult rewritePartitionDependencies(scf::ForOp &loop); -// Given a loop where the partitions' inputs and outputs have been fully -// rewritten to be reference semantic, partitiong the loop into a -// `ttg.warp_specialize` by duplicating the loop for each partition and -// rematerializing, as necessary, operations in the root partition. -LogicalResult partitionLoop(scf::ForOp loop); -} // namespace triton::gpu -} // namespace mlir - -#endif // TRITON_TRITONGPU_TRANSFORM_PIPELINE_WARPSPECIALIZATION_H_ diff --git a/third_party/mthreads/include/triton/Dialect/TritonInstrument/CMakeLists.txt b/third_party/mthreads/include/triton/Dialect/TritonInstrument/CMakeLists.txt deleted file mode 100644 index 9f57627c32..0000000000 --- a/third_party/mthreads/include/triton/Dialect/TritonInstrument/CMakeLists.txt +++ /dev/null @@ -1,2 +0,0 @@ -add_subdirectory(IR) -add_subdirectory(Transforms) diff --git a/third_party/mthreads/include/triton/Dialect/TritonInstrument/IR/CMakeLists.txt b/third_party/mthreads/include/triton/Dialect/TritonInstrument/IR/CMakeLists.txt deleted file mode 100644 index 2af09f9046..0000000000 --- a/third_party/mthreads/include/triton/Dialect/TritonInstrument/IR/CMakeLists.txt +++ /dev/null @@ -1,15 +0,0 @@ -set(MLIR_BINARY_DIR ${CMAKE_BINARY_DIR}) - -set(LLVM_TARGET_DEFINITIONS TritonInstrumentDialect.td) -mlir_tablegen(Dialect.h.inc -gen-dialect-decls -dialect=tti) -mlir_tablegen(Dialect.cpp.inc -gen-dialect-defs -dialect=tti) -add_mlir_doc(TritonInstrumentDialect TritonInstrumentDialect dialects/ -gen-dialect-doc) - -set(LLVM_TARGET_DEFINITIONS TritonInstrumentOps.td) -mlir_tablegen(Ops.h.inc -gen-op-decls) -mlir_tablegen(Ops.cpp.inc -gen-op-defs) -mlir_tablegen(OpsEnums.h.inc -gen-enum-decls) -mlir_tablegen(OpsEnums.cpp.inc -gen-enum-defs) -add_mlir_doc(TritonInstrumentOps TritonInstrumentOps dialects/ -gen-op-doc) - -add_public_tablegen_target(TritonInstrumentTableGen) diff --git a/third_party/mthreads/include/triton/Dialect/TritonInstrument/IR/Dialect.h b/third_party/mthreads/include/triton/Dialect/TritonInstrument/IR/Dialect.h deleted file mode 100644 index e0fcf61b44..0000000000 --- a/third_party/mthreads/include/triton/Dialect/TritonInstrument/IR/Dialect.h +++ /dev/null @@ -1,14 +0,0 @@ -#ifndef TRITON_DIALECT_TRITONINSTRUMENT_IR_DIALECT_H_ -#define TRITON_DIALECT_TRITONINSTRUMENT_IR_DIALECT_H_ - -// TritonInstrument depends on Triton and TritonGPU -#include "triton/Dialect/Triton/IR/Dialect.h" -#include "triton/Dialect/TritonGPU/IR/Dialect.h" - -#include "triton/Dialect/TritonInstrument/IR/OpsEnums.h.inc" - -#define GET_OP_CLASSES -#include "triton/Dialect/TritonInstrument/IR/Dialect.h.inc" -#include "triton/Dialect/TritonInstrument/IR/Ops.h.inc" - -#endif // TRITON_DIALECT_TRITONINSTRUMENT_IR_DIALECT_H_ diff --git a/third_party/mthreads/include/triton/Dialect/TritonInstrument/IR/TritonInstrument.md b/third_party/mthreads/include/triton/Dialect/TritonInstrument/IR/TritonInstrument.md deleted file mode 100644 index c7e05eef1d..0000000000 --- a/third_party/mthreads/include/triton/Dialect/TritonInstrument/IR/TritonInstrument.md +++ /dev/null @@ -1,86 +0,0 @@ -# Triton Instrument Dialect and Concurrency Sanitizer (ConSan) - -### Overview - -ConSan instruments Triton IR to detect illegal concurrent accesses to shared and Tensor Core memory under warp specialization. It tracks per-buffer visibility of reads and writes across threads, models barrier-based synchronization, and models commit-count–based synchronization (cp.async, wgmma). - -Auxiliary state is kept in distributed tensors and global scratch memory, with types created on-demand per warp-specialization partition. - -### Thread model - -- Base threads: 16 warp-specialization (WS) threads (allowing for up to 16 partitions). -- Peer classes: +16 Tensor Core (TC) threads and +16 TMA threads to model lack of ordering with base threads. -- Total logical threads: 48. Bitmasks are sized to the next power of two: 64. - -Indexing uses a logical thread id in [0, 48), with column vectors sized to 64 for layout convenience. - -## Auxiliary data structures - -All types are generated on-demand (per partition) based on: - -- B: number of tracked buffers (power-of-two padded) -- K: number of mbarriers (power-of-two padded) -- T_bits: 64 (bitmask width) -- T_commits: 16 (base threads; commit counters do not apply to TC/TMA helpers) - -“tensor” means a distributed Triton tensor; “scratch” means a pointer into global scratch memory. Shapes below are logical; actual encodings are partition-local blocked layouts. - -- buffers (tensor, ): Base pointers of all (sub)buffers per memory space -- barriers (tensor, ): Pointers of all mbarriers -- writeVisibility (scratch, ): Per-buffer bitmask. Bit i set ⇒ thread i can see latest completed write to that buffer -- readVisibility (scratch, ): Per-buffer, per-thread lanes. Each lane stores a 64-bit mask of other threads whose reads are visible to that lane’s thread -- writeTracking (scratch, ): Map buffers → barriers tracking writes (boolean stored in i8) -- readTracking (scratch, ): Map buffers → barriers tracking reads (bitmask of threads) -- barrierStates (scratch, ): Packed barrier metadata. Bit 0 stores the current phase, bits [1..8] the initial arrival count, bits [9..16] the current arrival count. The verifier checks underflow before updating, and flips the phase when the current count reaches zero. -- waiting (scratch, ): Per-barrier bitfield describing waiting threads. Each base thread gets two bits: bit (2 * thread + 0) is the waiting flag, bit (2 * thread + 1) stores the phase the thread is waiting on. -- outstandingCommits (scratch, ): Per-buffer, per-base-thread commit counters for cp.async and wgmma - -## Visibility and legality rules - -- Reads are legal iff the reading thread sees the most recent write to the buffer (writeVisibility). There can be only one write in-flight. -- Writes are legal iff the writing thread sees both all prior writes and all reads completed for that buffer. - -ConSan enforces these via two checks emitted before memory ops: - -- experimental_verify_write_visibility: “no one else is writing, or I can see the write” -- experimental_verify_read_visibility: “my read-visibility lane is a superset of the OR of all lanes” - -## Barrier-based synchronization - -ConSan separates “tracking” from “visibility transfer”: - -- At memory ops that are tracked by a barrier (loads/stores, some TMEM ops): - - experimental_set_read_visibility / experimental_set_write_visibility updates the appropriate visibility table for the current thread and buffer. - - experimental_track_visible_reads / experimental_track_visible_writes snapshots current per-buffer visibility into readTracking/writeTracking for the given barrier. -- At arrive/commit sites (e.g., tc commit, arrive on mbarrier): ConSan emits the track ops for both reads and writes. -- At waits: experimental_transfer_visible_reads / experimental_transfer_visible_writes propagates tracked visibility from the barrier back into the waiting thread’s visibility, and this transfer is repeated to peer threads (base, TMA, TC) to keep the three classes consistent. - -### Barrier phase/count tracking - -- experimental_init_barrier_state(barrier, count, barrierStates) initializes the per-barrier state with phase = 0 and both initial/current arrival counts = `count`. -- experimental_verify_barrier_arrive(barrier, count, barrierStates) checks that subtracting `count` from the current arrival count would not underflow. The codegen emits an assert if it would. -- experimental_update_barrier_state(barrier, count, barrierStates) applies the arrive: subtracts `count`, flips the phase when the count reaches zero, and reloads the current count from the initial count. - -### Deadlock detection - -ConSan records which phase each thread is waiting on: - -- experimental_set_waiting(barrier, baseThread, phase, barriers, waiting) sets the waiting flag for `baseThread` and stores the requested `phase`. The flag/phase bits share the waiting bitfield (two bits per base thread). -- experimental_check_all_active_waiting(activeMask, barriers, waiting, barrierStates) filters waiting threads to those whose stored phase matches the current barrier phase. If all active threads are waiting on matching phases, it raises a deadlock assert. -- experimental_clear_waiting(barrier, baseThread, barriers, waiting) clears the waiting bits for `baseThread`. Each wait clears its own state after the wait completes. - -## Commit-count–based synchronization - -Some hardware ops synchronize via “number of outstanding commits” rather than mbarriers. - -- Stage: experimental_stage_access_for_commit marks the current thread’s buffer lane with -1 (staged) in outstandingCommits[B x 16]. -- Commit: experimental_commit_accesses turns -1 into 1 and increments positive entries for the committing thread column. -- Wait (cp.async): experimental_clear_outstanding_commits_set_write(thread, commits, writeVisibility, N) clears entries with count > N for the current thread, and sets the writeVisibility bit for rows where any thread’s entry was cleared. -- Wait (wgmma): experimental_clear_outstanding_commits_set_read(thread, commits, readVisibility, N) clears entries with count > N for the current thread, and sets the readVisibility bit for rows where any thread’s entry was cleared. - -Legality checks for commit-count flows: - -- For writes to shared memory affected by cp.async: experimental_check_outstanding_commits(buffer, commits, "async_copy_global_to_shared") asserts the row for the buffer is all zeros (no pending writes), across all base-thread columns. -- For reads of wgmma operands in shared memory: experimental_check_outstanding_commits(buffer, commits, "warpgroup_mma operand read") asserts the row is all zeros (no pending reads). - -Note: The check op has no “thread” operand; it inspects the whole row for the buffer. diff --git a/third_party/mthreads/include/triton/Dialect/TritonInstrument/IR/TritonInstrumentAttrDefs.td b/third_party/mthreads/include/triton/Dialect/TritonInstrument/IR/TritonInstrumentAttrDefs.td deleted file mode 100644 index ab8702defb..0000000000 --- a/third_party/mthreads/include/triton/Dialect/TritonInstrument/IR/TritonInstrumentAttrDefs.td +++ /dev/null @@ -1,15 +0,0 @@ -#ifndef TRITONINSTRUMENT_ATTR_DEFS -#define TRITONINSTRUMENT_ATTR_DEFS - -include "mlir/IR/EnumAttr.td" - -def TT_MemTypeAttr : I32EnumAttr< - "MemType", "", - [ - I32EnumAttrCase<"SHARED_MEM", 0, "shared_mem">, - I32EnumAttrCase<"TENSOR_MEM", 1, "tensor_mem">, - ]> { - let cppNamespace = "::mlir::triton::instrument"; -} - -#endif // TRITONINSTRUMENT_ATTR_DEFS diff --git a/third_party/mthreads/include/triton/Dialect/TritonInstrument/IR/TritonInstrumentDialect.td b/third_party/mthreads/include/triton/Dialect/TritonInstrument/IR/TritonInstrumentDialect.td deleted file mode 100644 index 6a7f3eed62..0000000000 --- a/third_party/mthreads/include/triton/Dialect/TritonInstrument/IR/TritonInstrumentDialect.td +++ /dev/null @@ -1,11 +0,0 @@ -#ifndef TRITONINSTRUMENT_DIALECT -#define TRITONINSTRUMENT_DIALECT - -include "mlir/IR/OpBase.td" - -def TritonInstrument_Dialect : Dialect { - let name = "tti"; - let cppNamespace = "::mlir::triton::instrument"; -} - -#endif // TRITONINSTRUMENT_DIALECT diff --git a/third_party/mthreads/include/triton/Dialect/TritonInstrument/Transforms/CMakeLists.txt b/third_party/mthreads/include/triton/Dialect/TritonInstrument/Transforms/CMakeLists.txt deleted file mode 100644 index 672815ac4b..0000000000 --- a/third_party/mthreads/include/triton/Dialect/TritonInstrument/Transforms/CMakeLists.txt +++ /dev/null @@ -1,3 +0,0 @@ -set(LLVM_TARGET_DEFINITIONS Passes.td) -mlir_tablegen(Passes.h.inc -gen-pass-decls -name TritonInstrument) -add_public_tablegen_target(TritonInstrumentTransformsIncGen) diff --git a/third_party/mthreads/include/triton/Dialect/TritonInstrument/Transforms/Passes.h b/third_party/mthreads/include/triton/Dialect/TritonInstrument/Transforms/Passes.h deleted file mode 100644 index c96c618e68..0000000000 --- a/third_party/mthreads/include/triton/Dialect/TritonInstrument/Transforms/Passes.h +++ /dev/null @@ -1,22 +0,0 @@ -#ifndef TRITON_DIALECT_TRITONINSTRUMENT_TRANSFORMS_PASSES_H_ -#define TRITON_DIALECT_TRITONINSTRUMENT_TRANSFORMS_PASSES_H_ - -#include "mlir/Pass/Pass.h" -#include "triton/Dialect/TritonInstrument/IR/Dialect.h" - -namespace mlir { -namespace triton { -namespace instrument { - -// Generate the pass class declarations. -#define GEN_PASS_DECL -#include "triton/Dialect/TritonInstrument/Transforms/Passes.h.inc" - -/// Generate the code for registering passes. -#define GEN_PASS_REGISTRATION -#include "triton/Dialect/TritonInstrument/Transforms/Passes.h.inc" - -} // namespace instrument -} // namespace triton -} // namespace mlir -#endif diff --git a/third_party/mthreads/include/triton/Dialect/TritonInstrument/Transforms/Passes.td b/third_party/mthreads/include/triton/Dialect/TritonInstrument/Transforms/Passes.td deleted file mode 100644 index cfd860e991..0000000000 --- a/third_party/mthreads/include/triton/Dialect/TritonInstrument/Transforms/Passes.td +++ /dev/null @@ -1,16 +0,0 @@ -#ifndef TRITONINSTRUMENT_PASSES -#define TRITONINSTRUMENT_PASSES - -include "mlir/Pass/PassBase.td" - -def TritonInstrumentConcurrencySanitizer: Pass<"tritoninstrument-concurrency-sanitizer", "mlir::ModuleOp"> { - let summary = "Add runtime verification of asynchronous operations"; - - let description = "Instrument the program with runtime verification of asynchronous operations."; - - let dependentDialects = ["mlir::triton::gpu::TritonGPUDialect", - "mlir::triton::TritonDialect", - "mlir::triton::instrument::TritonInstrumentDialect"]; -} - -#endif // TRITON_INSTRUMENT_PASSES diff --git a/third_party/mthreads/include/triton/Dialect/TritonNvidiaGPU/CMakeLists.txt b/third_party/mthreads/include/triton/Dialect/TritonNvidiaGPU/CMakeLists.txt deleted file mode 100644 index 9f57627c32..0000000000 --- a/third_party/mthreads/include/triton/Dialect/TritonNvidiaGPU/CMakeLists.txt +++ /dev/null @@ -1,2 +0,0 @@ -add_subdirectory(IR) -add_subdirectory(Transforms) diff --git a/third_party/mthreads/include/triton/Dialect/TritonNvidiaGPU/IR/CMakeLists.txt b/third_party/mthreads/include/triton/Dialect/TritonNvidiaGPU/IR/CMakeLists.txt deleted file mode 100644 index b93aad2ba0..0000000000 --- a/third_party/mthreads/include/triton/Dialect/TritonNvidiaGPU/IR/CMakeLists.txt +++ /dev/null @@ -1,27 +0,0 @@ -set(MLIR_BINARY_DIR ${CMAKE_BINARY_DIR}) - -set(LLVM_TARGET_DEFINITIONS TritonNvidiaGPUOps.td) -mlir_tablegen(Dialect.h.inc -gen-dialect-decls -dialect=ttng) -mlir_tablegen(Dialect.cpp.inc -gen-dialect-defs -dialect=ttng) -mlir_tablegen(Ops.h.inc -gen-op-decls) -mlir_tablegen(Ops.cpp.inc -gen-op-defs) -add_mlir_doc(TritonNvidiaGPUDialect TritonNvidiaGPUDialect dialects/ -gen-dialect-doc) -add_mlir_doc(TritonNvidiaGPUOps TritonNvidiaGPUOps dialects/ -gen-op-doc) -add_public_tablegen_target(TritonNvidiaGPUTableGen) - -set(LLVM_TARGET_DEFINITIONS TritonNvidiaGPUTypes.td) -mlir_tablegen(Types.h.inc -gen-typedef-decls) -mlir_tablegen(Types.cpp.inc -gen-typedef-defs) -add_public_tablegen_target(TritonNvidiaGPUTypesIncGen) - -set(LLVM_TARGET_DEFINITIONS TritonNvidiaGPUAttrDefs.td) -mlir_tablegen(TritonNvidiaGPUAttrDefs.h.inc -gen-attrdef-decls) -mlir_tablegen(TritonNvidiaGPUAttrDefs.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(TritonNvidiaGPUAttrDefsIncGen) - -set(LLVM_TARGET_DEFINITIONS TritonNvidiaGPUOpInterfaces.td) -mlir_tablegen(TritonNvidiaGPUOpInterfaces.h.inc -gen-op-interface-decls) -mlir_tablegen(TritonNvidiaGPUOpInterfaces.cpp.inc -gen-op-interface-defs) -add_public_tablegen_target(TritonNvidiaGPUOpInterfacesIncGen) diff --git a/third_party/mthreads/include/triton/Dialect/TritonNvidiaGPU/IR/TensorMemoryUtils.h b/third_party/mthreads/include/triton/Dialect/TritonNvidiaGPU/IR/TensorMemoryUtils.h deleted file mode 100644 index 3ae002a597..0000000000 --- a/third_party/mthreads/include/triton/Dialect/TritonNvidiaGPU/IR/TensorMemoryUtils.h +++ /dev/null @@ -1,37 +0,0 @@ -#ifndef TRITON_DIALECT_TRITONNVIDIAGPU_IR_TENSORMEMORYUTILS_H_ -#define TRITON_DIALECT_TRITONNVIDIAGPU_IR_TENSORMEMORYUTILS_H_ - -#include "mlir/IR/BuiltinTypes.h" -#include "triton/Dialect/TritonNvidiaGPU/IR/Dialect.h" -#include "triton/Tools/LinearLayout.h" - -#include -#include -#include - -namespace mlir::triton::nvidia_gpu { - -// Get the maximum number of registers per thread based on the context. This is -// by default 256, but it can be overridden by `ttg.maxnreg` set on the module -// or a contextual register limit set by the compiler on partitions. -int getContextualMaxNReg(Operation *op); -struct TMemLdStEncodingInfo { - TMemAccessAtom atom; - LinearLayout reps; - ColumnAction perm; - int numRegsPerMessage; - std::optional secondHalfOffset; - std::optional broadcast = std::nullopt; - bool unpacked = false; - unsigned vec = 1; - bool padding = false; -}; - -FailureOr -computeTMemLdStEncodingInfo(RankedTensorType regTy, gpu::MemDescType memTy, - int maxnreg, - std::function emitError = {}); - -} // namespace mlir::triton::nvidia_gpu - -#endif // TRITON_DIALECT_TRITONNVIDIAGPU_IR_TENSORMEMORYUTILS_H_ diff --git a/third_party/mthreads/include/triton/Dialect/TritonNvidiaGPU/Transforms/CMakeLists.txt b/third_party/mthreads/include/triton/Dialect/TritonNvidiaGPU/Transforms/CMakeLists.txt deleted file mode 100644 index d4b5c097f4..0000000000 --- a/third_party/mthreads/include/triton/Dialect/TritonNvidiaGPU/Transforms/CMakeLists.txt +++ /dev/null @@ -1,3 +0,0 @@ -set(LLVM_TARGET_DEFINITIONS Passes.td) -mlir_tablegen(Passes.h.inc -gen-pass-decls -name TritonNvidiaGPU) -add_public_tablegen_target(TritonNvidiaGPUTransformsIncGen) diff --git a/third_party/mthreads/include/triton/Dialect/TritonNvidiaGPU/Transforms/Passes.h b/third_party/mthreads/include/triton/Dialect/TritonNvidiaGPU/Transforms/Passes.h deleted file mode 100644 index b11a3f653e..0000000000 --- a/third_party/mthreads/include/triton/Dialect/TritonNvidiaGPU/Transforms/Passes.h +++ /dev/null @@ -1,46 +0,0 @@ -/* - * Copyright (c) 2023 NVIDIA Corporation & Affiliates. All rights reserved. - * - * Permission is hereby granted, free of charge, to any person obtaining - * a copy of this software and associated documentation files - * (the "Software"), to deal in the Software without restriction, - * including without limitation the rights to use, copy, modify, merge, - * publish, distribute, sublicense, and/or sell copies of the Software, - * and to permit persons to whom the Software is furnished to do so, - * subject to the following conditions: - * - * The above copyright notice and this permission notice shall be - * included in all copies or substantial portions of the Software. - * - * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, - * EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF - * MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. - * IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY - * CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, - * TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE - * SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. - */ - -#ifndef TRITON_DIALECT_TRITONNVIDIAGPU_TRANSFORMS_PASSES_H_ -#define TRITON_DIALECT_TRITONNVIDIAGPU_TRANSFORMS_PASSES_H_ - -#include "mlir/Pass/Pass.h" - -namespace mlir { -namespace triton { -namespace nvidia_gpu { - -std::unique_ptr createTritonNvidiaGPUPlanCTAPass(); - -#define GEN_PASS_DECL -#include "triton/Dialect/TritonNvidiaGPU/Transforms/Passes.h.inc" - -/// Generate the code for registering passes. -#define GEN_PASS_REGISTRATION -#include "triton/Dialect/TritonNvidiaGPU/Transforms/Passes.h.inc" - -} // namespace nvidia_gpu -} // namespace triton -} // namespace mlir - -#endif // TRITON_DIALECT_TRITONNVIDIAGPU_TRANSFORMS_PASSES_H_ diff --git a/third_party/mthreads/include/triton/Dialect/TritonNvidiaGPU/Transforms/Passes.td b/third_party/mthreads/include/triton/Dialect/TritonNvidiaGPU/Transforms/Passes.td deleted file mode 100644 index a41b2e8914..0000000000 --- a/third_party/mthreads/include/triton/Dialect/TritonNvidiaGPU/Transforms/Passes.td +++ /dev/null @@ -1,187 +0,0 @@ -// Copyright (c) 2023 NVIDIA Corporation & Affiliates. All rights reserved. -// -// Permission is hereby granted, free of charge, to any person obtaining -// a copy of this software and associated documentation files -// (the "Software"), to deal in the Software without restriction, -// including without limitation the rights to use, copy, modify, merge, -// publish, distribute, sublicense, and/or sell copies of the Software, -// and to permit persons to whom the Software is furnished to do so, -// subject to the following conditions: -// -// The above copyright notice and this permission notice shall be -// included in all copies or substantial portions of the Software. -// -// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, -// EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF -// MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. -// IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY -// CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, -// TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE -// SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. - -#ifndef TRITONNVIDIAGPU_PASSES -#define TRITONNVIDIAGPU_PASSES - -include "mlir/Pass/PassBase.td" - -def TritonGPUPlanCTAPass : Pass<"triton-nvidia-gpu-plan-cta", "mlir::ModuleOp"> { - let summary = "plan CTA"; - - let description = [{ - This pass computes and applies "optimized" CTA tilings to DotOp, ReduceOp - and StoreLikeOps operations. - }]; - - let constructor = "mlir::triton::nvidia_gpu::createTritonNvidiaGPUPlanCTAPass()"; - - let dependentDialects = [ - "mlir::triton::gpu::TritonGPUDialect", - "mlir::triton::nvidia_gpu::TritonNvidiaGPUDialect" - ]; -} - -def TritonGPUFenceInsertion : Pass<"triton-nvidia-gpu-fence-insertion", "mlir::ModuleOp"> { - let summary = "Insert fences across generic and async proxy."; - - let description = [{ - This pass is to insert memory fences to ensure that memory operations are - properly ordered across generic and async operations. - This pass inserts fences at optimized location. - There is a pass later to handle all the functional requirements - }]; - - let dependentDialects = [ - "mlir::triton::gpu::TritonGPUDialect", - "mlir::triton::nvidia_gpu::TritonNvidiaGPUDialect" - ]; - - let options = [ - Option<"computeCapability", "compute-capability", - "int32_t", /*default*/"90", - "device compute capability"> - ]; -} - -def TritonGPUProxyFenceInsertion : Pass<"triton-nvidia-gpu-proxy-fence-insertion", "mlir::ModuleOp"> { - let summary = "Insert fences across generic and async proxy"; - - let description = [{ - This pass is to insert memory fences to ensure that memory operations are - properly ordered across generic and async operations. - }]; - - let dependentDialects = [ - "mlir::triton::gpu::TritonGPUDialect", - "mlir::triton::nvidia_gpu::TritonNvidiaGPUDialect" - ]; - - let options = [ - Option<"computeCapability", "compute-capability", - "int32_t", /*default*/"90", - "device compute capability"> - ]; -} - -def TritonNvidiaGPUTMALoweringPass : Pass<"triton-nvidia-tma-lowering", "mlir::ModuleOp"> { - let summary = "lower to TMA load/store operations"; - - let description = [{ - Lower Triton descriptor load to TMA load/store operations in TritonNvidiaGPUDialect. - }]; - - let dependentDialects = [ - "mlir::triton::nvidia_gpu::TritonNvidiaGPUDialect" - ]; -} - -def TritonTensorMemoryAllocationPass : Pass<"triton-tensor-memory-allocation", "mlir::ModuleOp"> { - let summary = "Assign tensor memory allocation"; - - let description = [{ - Decide on tensor memory allocation and assign attributes to each allocation. - }]; - - let dependentDialects = [ - "mlir::triton::nvidia_gpu::TritonNvidiaGPUDialect" - ]; -} - -def TritonNvidiaGPUMMALoweringPass : Pass<"triton-nvidia-mma-lowering", "mlir::ModuleOp"> { - let summary = "lower mma operations if needed"; - - let description = [{ - Lower MMA ops to prepare for conversion to LLVM. - }]; - - let dependentDialects = [ - "mlir::triton::nvidia_gpu::TritonNvidiaGPUDialect" - ]; -} - -def TritonNvidiaGPUPromoteLHSToTMemPass : Pass<"tritongpu-promote-lhs-to-tmem", "mlir::ModuleOp"> { - let summary = "Promote LHS operand of MMAv5 op to Tensor Memory"; - - let description = [{ - Promote LHS operand of MMAv5 op to Tensor Memory. - }]; - - let dependentDialects = ["mlir::triton::gpu::TritonGPUDialect", - "mlir::triton::nvidia_gpu::TritonNvidiaGPUDialect", - "mlir::triton::TritonDialect"]; -} - -def TritonNvidiaGPUOptimizeDescriptorEncodingPass : Pass<"triton-nvidia-optimize-descriptor-encoding", "mlir::ModuleOp"> { - let summary = "Set encodings on tensor descriptor types"; - - let description = [{ - Set shared memory encoding on tensor descriptors, which decides the swizzling mode and message size of the tma descriptor. - }]; - - let dependentDialects = ["mlir::triton::gpu::TritonGPUDialect", - "mlir::triton::nvidia_gpu::TritonNvidiaGPUDialect", - "mlir::triton::TritonDialect"]; -} - -def TritonNvidiaGPUOptimizeTMemLayoutsPass : Pass<"triton-nvidia-optimize-tmem-layouts", "mlir::ModuleOp"> { - let summary = "Optimize TMEM layouts."; - - let description = [{ - Optimize TMEM layouts by selecting a layouts to enable better subtiling, - reduction performance, etc. - }]; - - let dependentDialects = ["mlir::triton::gpu::TritonGPUDialect", - "mlir::triton::nvidia_gpu::TritonNvidiaGPUDialect", - "mlir::triton::TritonDialect"]; -} - -def TritonNvidiaGPUInterleaveTMemPass : Pass<"triton-nvidia-interleave-tmem", "mlir::ModuleOp"> { - let summary = "Interleave TMEM loads/stores."; - - let description = [{ - The `triton-nvidia-interleave-tmem` pass attempts to sink TMEM loads and - hoist TMEM stores, and potentially interleave them, to reduce register - pressure. - }]; -} - -def TritonNvidiaGPURemoveTMEMTokensPass : Pass<"triton-nvidia-gpu-remove-tmem-tokens", "mlir::ModuleOp"> { - let summary = "remove TMEM tokens"; - - let description = [{ - The `triton-nvidia-gpu-remove-tmem-tokens` pass removes TMEM memory - dependency tokens from the IR, after they are no longer needed. - }]; -} - -def TritonNvidiaGPUCheckMatmulTwoCTAPass : Pass<"triton-nvidia-check-matmul-two-cta", "mlir::ModuleOp"> { - let summary = "Verify consistent two_ctas usage across matmuls"; - - let description = [{ - Inspect all matmul operations and ensure they agree on the `two_ctas` - setting. Propagate the chosen value to the module so later lowering steps - can access it. Compilation fails if mixed configurations are detected. - }]; -} - -#endif diff --git a/third_party/mthreads/include/triton/Target/CMakeLists.txt b/third_party/mthreads/include/triton/Target/CMakeLists.txt deleted file mode 100644 index 39d31dc9b5..0000000000 --- a/third_party/mthreads/include/triton/Target/CMakeLists.txt +++ /dev/null @@ -1 +0,0 @@ -add_subdirectory(LLVMIR) diff --git a/third_party/mthreads/include/triton/Target/LLVMIR/CMakeLists.txt b/third_party/mthreads/include/triton/Target/LLVMIR/CMakeLists.txt deleted file mode 100644 index 1f6c1b3511..0000000000 --- a/third_party/mthreads/include/triton/Target/LLVMIR/CMakeLists.txt +++ /dev/null @@ -1,3 +0,0 @@ -set(LLVM_TARGET_DEFINITIONS Passes.td) -mlir_tablegen(Passes.h.inc -gen-pass-decls -name LLVMIR) -add_public_tablegen_target(LLVMIRIncGen) diff --git a/third_party/mthreads/include/triton/Target/LLVMIR/Passes.h b/third_party/mthreads/include/triton/Target/LLVMIR/Passes.h deleted file mode 100644 index 87da907e14..0000000000 --- a/third_party/mthreads/include/triton/Target/LLVMIR/Passes.h +++ /dev/null @@ -1,18 +0,0 @@ -#ifndef TRITON_TARGET_LLVM_IR_PASSES_H -#define TRITON_TARGET_LLVM_IR_PASSES_H - -#include "mlir/Pass/Pass.h" - -namespace mlir { - -// Generate the pass class declarations. -#define GEN_PASS_DECL -#include "triton/Target/LLVMIR/Passes.h.inc" - -// Generate the code for registering conversion passes. -#define GEN_PASS_REGISTRATION -#include "triton/Target/LLVMIR/Passes.h.inc" - -} // namespace mlir - -#endif // TRITON_TARGET_LLVM_IR_PASSES_H diff --git a/third_party/mthreads/include/triton/Target/LLVMIR/Passes.td b/third_party/mthreads/include/triton/Target/LLVMIR/Passes.td deleted file mode 100644 index 854d753342..0000000000 --- a/third_party/mthreads/include/triton/Target/LLVMIR/Passes.td +++ /dev/null @@ -1,21 +0,0 @@ -#ifndef TRITON_TARGET_LLVMIR_PASSES -#define TRITON_TARGET_LLVMIR_PASSES - -include "mlir/Pass/PassBase.td" - -def LLVMDIScope: Pass<"enable-line-info", "mlir::ModuleOp"> { - let summary = "Materialize LLVM line info"; - let description = [{ - This pass materializes line mapping information for LLVM IR dialect operations. - }]; -} - -def LLVMDILocalVariable: Pass<"extract-variable-info", "mlir::ModuleOp"> { - let summary = "Pull out source variable info from Location to DILocalVariable"; - let description = [{ - This pass pulled out source vararible's debuginfo from LLVM IR dialect's Location - into LLVM's DILocalVariable and fused it into previous Location so it can be passed to LLVM IR later in debugging mode. - }]; -} - -#endif diff --git a/third_party/mthreads/include/triton/Tools/GenericSwizzling.h b/third_party/mthreads/include/triton/Tools/GenericSwizzling.h deleted file mode 100644 index e1b3b3e2cc..0000000000 --- a/third_party/mthreads/include/triton/Tools/GenericSwizzling.h +++ /dev/null @@ -1,56 +0,0 @@ -#ifndef TRITON_GENERIC_SWIZZLING_H -#define TRITON_GENERIC_SWIZZLING_H - -#include "llvm/ADT/ArrayRef.h" -#include "llvm/ADT/SmallVector.h" -#include -#include - -namespace mlir::triton { -class LinearLayout; -class TargetInfoBase; -} // namespace mlir::triton - -namespace mlir::triton::gpu { -// Store the lane indices that are used in the contiguous part -// of an operation and in the address part. -// The laneAddr part just represents the indices used in one wavefront -// For now we just represent tiles with full vectorisation, meaning -// ld.shared.b32.v4/st.shared.b32.v4 -// ldmatrix.v4 / stmatrix.v4 -// ldmatrix.trans.v4 / stmatrix.trans.v4 -struct LocalMemOpTile { - // If laneContig.size() < log2(128/bitwidth), we assume that - // the first log2(128/bitwidth) - laneContig.size() bases are registers - llvm::SmallVector laneContig; - // If laneAddr.size() < 3, we assume that the first - // 3 - laneAddr.size() bases are registers - llvm::SmallVector laneAddr; -}; - -// Given a set of possible instructions given by -// targetInfo.laneIdTiles(bitwidth) returns the optimal swizzling given these -// instructions and a pair of indices into the ldStTiles that's needed to lower -// this swizzling -std::pair> -optimalSwizzling(const LinearLayout &src, const LinearLayout &dst, - llvm::ArrayRef srcTiles, - llvm::ArrayRef dstTiles, int32_t bitwidth); - -LinearLayout optimalSwizzlingLdSt(const LinearLayout &src, - const LinearLayout &dst, int32_t bitwidth); - -std::pair bankConflictsLdSt(const LinearLayout &src, - const LinearLayout &dst, - const LinearLayout &smem, - int32_t bitwidth); - -int bankConflictsMemDesc(const LinearLayout ®, const LinearLayout &smem, - int32_t bitwidth); - -std::pair bankConflicts(llvm::ArrayRef tileSrc, - llvm::ArrayRef tileDst, - const LinearLayout &smem); -} // namespace mlir::triton::gpu - -#endif // TRITON_GENERIC_SWIZZLING_H diff --git a/third_party/mthreads/include/triton/Tools/LayoutUtils.h b/third_party/mthreads/include/triton/Tools/LayoutUtils.h deleted file mode 100644 index 7ea612fb02..0000000000 --- a/third_party/mthreads/include/triton/Tools/LayoutUtils.h +++ /dev/null @@ -1,190 +0,0 @@ -#ifndef TRITON_TOOLS_LAYOUTUTILS_H -#define TRITON_TOOLS_LAYOUTUTILS_H - -#include "triton/Tools/LinearLayout.h" - -namespace mlir::triton { -// Is the sublayout defined from dimNames to dimNames the identity? -// In particular, is the input and output size in these dimensions -// the same, and are the bases the identity? -bool squareSublayoutIsIdentity(const LinearLayout &ll, - ArrayRef dimNames); - -// For each output dimension d, ensure that the layout's output size (i.e., its -// codomain) does not exceed shape[d]. Do this without changing the size of the -// layout's inputs (i.e., leave its domain unchanged). -// -// This function is invariant to the order of the layout's input and output -// dimensions. -// -// We achieve this by setting the largest value in each output dimension d to 0 -// because bases that map to a location larger than shape[d] -// effectively duplicate along that dimension. For example, consider a layout -// with an output dimension size of 32, and we call ensureLayoutNotLargerThan to -// shrink the output dimension size to 8: -// -// L(register=1) = 8 -// L(register=2) = 4 -// L(register=4) = 1 -// L(lane=1) = 2 -// L(lane=2) = 16 -// -// In the first step, we shrink the output dimension size to 16 by setting -// L(lane=2) to 0: -// -// L(register=1) = 8 -// L(register=2) = 4 -// L(register=4) = 1 -// L(lane=1) = 2 -// L(lane=2) = 0 -// -// This means that lane=2 has the same data as lane=0. -// -// Now the output dimension of this layout has a size of 16, which is still -// larger than 8. We find the current largest value in the output dimension, -// which is L(register=1) = 8, and we set L(register=1) to 0: -// -// L(register=1) = 0 -// L(register=2) = 4 -// L(register=4) = 1 -// L(lane=1) = 2 -// L(lane=2) = 0 -// -// Now the output dimension of this layout has a size of 8, which is the desired -// size. Note that this method works only because the bases are powers of two, -// which is the case for DistributedLayouts If broadcastRegisters is false, we -// remove any register that's larger than the desired shape. In the example -// above we would have -// L(register=1) = 4 -// L(register=2) = 1 -// L(lane=1) = 2 -// L(lane=2) = 0 -LinearLayout -ensureLayoutNotLargerThan(const LinearLayout &layout, - const llvm::SmallDenseMap &shape, - bool broadcastRegisters = true); - -// For each out-dim d, ensure the layout's out-size (i.e. its codomain) is no -// smaller than shape[d]. Do this by increasing the size of the layout's inputs -// along its most-minor dimension ("register" for register layouts, "offset" for -// shared layouts). -// -// This function is invariant to the order of the layout's input dimensions, but -// it cares about the order of the output dims, which should be minor-to-major. -LinearLayout ensureLayoutNotSmallerThan( - const LinearLayout &layout, - const llvm::SmallDenseMap &shape); - -inline LinearLayout -ensureLayoutNotSmallerThan(const LinearLayout &layout, - const llvm::ArrayRef dimNames, - const llvm::ArrayRef shape) { - llvm::SmallDenseMap namedDims; - for (auto [dimName, length] : llvm::zip_equal(dimNames, shape)) - namedDims[dimName] = length; - assert(namedDims.size() == shape.size() && "duplicate dimension names given"); - return ensureLayoutNotSmallerThan(layout, namedDims); -} - -// Return a vector of the standard out dimension names for tensor layouts. These -// are "dim0", "dim1", etc. -SmallVector standardOutDimNames(MLIRContext *ctx, int rank); - -// Return a vector of the standard out dimension name/value pairs, i.e. -// ("dim0", dstShape[0]), ("dim1", dstShape[1]), etc. -SmallVector> -standardOutDimPairs(MLIRContext *ctx, ArrayRef dstShape); - -// Return an identity mapping from `inDimName` to the standard out dimensions, -// with the dimensions sized according to the shape. The bases are sorted -// according to `order`, with the most minor dimension first. -LinearLayout identityStandardND(StringAttr inDimName, ArrayRef shape, - ArrayRef order); - -// Return a layout with the same in/out dimensions as `layout` but with all -// bases set to 0. -LinearLayout zerosLike(const LinearLayout &layout); - -// For a layout A with A.hasInDim(kReg), find a permutation of registers action -// such that action.apply(A) may be divisible by B -// It's not always true that the action returned by this function will -// allow us to divideLeft (resp. divideRight), but it is true that if it if -// there exists one, it is the one returned by this function. -std::optional regPermForDivide(const LinearLayout &A, - const LinearLayout &B, bool left); - -// For a layout A with A.hasInDim(kReg), find a permutation of registers action -// such that action.apply(A) has the broadcasted registers removed -ColumnAction actionRemoveBroadcastedRegs(const LinearLayout &layout); - -std::pair -actionAdditiveStrides(const LinearLayout &layout, const LinearLayout addrLayout, - uint64_t maskSpanOffsets); - -// For a layout A with A.hasInDim(kReg), repeat the values so that they have -// the same broadcasting as layout -SmallVector broadcastAs(const SmallVector &values, - const LinearLayout &layout); - -// Compute the supremum of two lists. -// Error out if the supremum does not exist (e.g. [a, b] and [b, a]). -// If the supremum is not unique, we return the first list first -// (e.g. [a, b], [a, c] -> [a, b, c]). -SmallVector supremum(const SmallVector &x, - const SmallVector &y); - -// Return a new layout reshaped to the given shape. -LinearLayout reshapeLayout(MLIRContext *ctx, LinearLayout layout, - ArrayRef shape); - -// Return a new layout with the dimensions transposed according to the given -// order. -LinearLayout transposeLinearLayout(LinearLayout layout, ArrayRef order); - -// Given a distributed into shmem layout, return the largest vectorisation -// that can be used to lower the layout via ld/st. -std::pair -largestVectorisation(MLIRContext *ctx, const LinearLayout &cvt, int bitwidth, - std::optional maybeMaxVecElems = std::nullopt); - -// Close cousin of doing zerosLike(tile) * divideLeft(cvt, tile) -// This one is a tad more general in the sense that it allows to divide -// cvt: -// - register=1 -> (0, 1) -// register=2 -> (8, 0) -// register=4 -> (0, 8) -// register=8 -> (0, 16) -// register=16 -> (0, 32) -// register=32 -> (0, 64) -// register=64 -> (16, 0) -// - lane=1 -> (0, 2) -// lane=2 -> (0, 4) -// lane=4 -> (1, 0) -// lane=8 -> (2, 0) -// lane=16 -> (4, 0) -// - warp=1 -> (32, 0) -// warp=2 -> (64, 0) -// - block is a size 1 dimension -// where out dims are: [row (size 128), col (size 128)] -// tile: -// - register=1 -> (0, 1) -// register=2 -> (8, 0) -// - lane=1 -> (0, 2) -// lane=2 -> (0, 4) -// lane=4 -> (1, 0) -// lane=8 -> (2, 0) -// lane=16 -> (4, 0) -// - warp=1 -> (32, 0) -// warp=2 -> (64, 0) -// where out dims are: [row (size 128), col (size 8)] -// which would not be possible to lower via the divideLeft approach as we -// cannot divide by the tile given the `register=64 -> (16, 0)` basis. -std::optional getReps(const LinearLayout &cvt, - const LinearLayout &tile); - -// Given a layout mapping onto dim0..dimn, remove a dimension `dim` -// and rename the rest as dim0..dimn-1 -LinearLayout removeStandardDim(const LinearLayout &layout, int dim); -} // namespace mlir::triton - -#endif // TRITON_TOOLS_LAYOUTUTILS_H diff --git a/third_party/mthreads/include/triton/Tools/StrUtil.h b/third_party/mthreads/include/triton/Tools/StrUtil.h deleted file mode 100644 index 8b59f7d2b3..0000000000 --- a/third_party/mthreads/include/triton/Tools/StrUtil.h +++ /dev/null @@ -1,54 +0,0 @@ -#include -#include - -#include "llvm/ADT/StringRef.h" -#include "llvm/Support/raw_ostream.h" - -namespace mlir::triton { - -// Better version of llvm::join. This one works when T is an integer or any -// other type which defines operator<<(raw_ostream). -template -std::string join(C &&container, llvm::StringRef sep = ", ") { - std::string ret; - llvm::raw_string_ostream s(ret); - for (const auto &elem : container) { - if (!ret.empty()) - s << sep; - s << elem; - } - return ret; -} - -// Joins a container of elements into a string, using `sep` as a separator. -// -// fn is called to transform each element of the container before it's added to -// the string. fn must have one of the following two signatures. -// -// - void fn(llvm::raw_ostream&, E), where E is the element type of the -// container, or -// - T fn(E), where T is a type which can be passed to -// raw_ostream::operator<<. -// -template -std::string join(C &&container, llvm::StringRef sep, Fn &&fn) { - std::string ret; - llvm::raw_string_ostream s(ret); - for (const auto &elem : container) { - if (!ret.empty()) - s << sep; - - if constexpr (std::is_invocable_v) { - static_assert( - std::is_void_v< - std::invoke_result_t>); - fn(s, elem); - } else { - s << fn(elem); - } - } - return ret; -} - -} // namespace mlir::triton diff --git a/third_party/mthreads/lib/Analysis/CMakeLists.txt b/third_party/mthreads/lib/Analysis/CMakeLists.txt deleted file mode 100644 index f7bd302ded..0000000000 --- a/third_party/mthreads/lib/Analysis/CMakeLists.txt +++ /dev/null @@ -1,23 +0,0 @@ -add_triton_library(TritonAnalysis - AxisInfo.cpp - Allocation.cpp - BufferRegion.cpp - Membar.cpp - Alias.cpp - Utility.cpp - - DEPENDS - TritonTableGen - TritonGPUTableGen - TritonGPUAttrDefsIncGen - TritonGPUTypeInterfacesIncGen - TritonGPUOpInterfacesIncGen - - LINK_LIBS PUBLIC - MLIRAnalysis - MLIRLLVMDialect - TritonIR - TritonGPUIR - GluonIR - TritonNvidiaGPUIR -) diff --git a/third_party/mthreads/lib/CMakeLists.txt b/third_party/mthreads/lib/CMakeLists.txt deleted file mode 100644 index c58b7fa0a3..0000000000 --- a/third_party/mthreads/lib/CMakeLists.txt +++ /dev/null @@ -1,5 +0,0 @@ -add_subdirectory(Analysis) -add_subdirectory(Conversion) -add_subdirectory(Dialect) -add_subdirectory(Target) -add_subdirectory(Tools) diff --git a/third_party/mthreads/lib/Conversion/CMakeLists.txt b/third_party/mthreads/lib/Conversion/CMakeLists.txt deleted file mode 100644 index 84aba4f3d2..0000000000 --- a/third_party/mthreads/lib/Conversion/CMakeLists.txt +++ /dev/null @@ -1,3 +0,0 @@ -add_subdirectory(TritonToTritonGPU) -add_subdirectory(TritonGPUToLLVM) -add_subdirectory(TritonInstrumentToLLVM) diff --git a/third_party/mthreads/lib/Conversion/TritonGPUToLLVM/AllocateSharedMemory.cpp b/third_party/mthreads/lib/Conversion/TritonGPUToLLVM/AllocateSharedMemory.cpp deleted file mode 100644 index 0448fbc73a..0000000000 --- a/third_party/mthreads/lib/Conversion/TritonGPUToLLVM/AllocateSharedMemory.cpp +++ /dev/null @@ -1,27 +0,0 @@ -#include "triton/Analysis/Allocation.h" -#include "triton/Analysis/Utility.h" -#include "triton/Conversion/TritonGPUToLLVM/AllocateSharedMemoryUtility.h" -#include "triton/Conversion/TritonGPUToLLVM/Passes.h" -#include "triton/Dialect/Triton/IR/Dialect.h" -#include "triton/Dialect/TritonGPU/IR/Dialect.h" - -using namespace mlir; -using namespace mlir::triton; - -namespace mlir::triton::gpu { -#define GEN_PASS_DEF_ALLOCATESHAREDMEMORY -#include "triton/Conversion/TritonGPUToLLVM/Passes.h.inc" -} // namespace mlir::triton::gpu - -namespace { -struct AllocateSharedMemory - : public mlir::triton::gpu::impl::AllocateSharedMemoryBase< - AllocateSharedMemory> { - void runOnOperation() override { - ModuleOp mod = getOperation(); - ModuleAllocation allocation(mod); - - mlir::triton::gpu::attachAllocationSizeAndOffsetAttr(mod, allocation); - } -}; -} // namespace diff --git a/third_party/mthreads/lib/Conversion/TritonGPUToLLVM/AllocateSharedMemoryUtility.cpp b/third_party/mthreads/lib/Conversion/TritonGPUToLLVM/AllocateSharedMemoryUtility.cpp deleted file mode 100644 index 24e90a2460..0000000000 --- a/third_party/mthreads/lib/Conversion/TritonGPUToLLVM/AllocateSharedMemoryUtility.cpp +++ /dev/null @@ -1,34 +0,0 @@ -#include "triton/Conversion/TritonGPUToLLVM/AllocateSharedMemoryUtility.h" - -namespace mlir::triton::gpu { - -void attachAllocationSizeAndOffsetAttr(ModuleOp mod, - ModuleAllocation &allocation) { - MLIRContext *ctx = mod.getContext(); - - mod.walk([&](FunctionOpInterface funcOp) { - auto *funcAllocation = allocation.getFuncData(funcOp); - funcOp.walk([&](Operation *op) { - auto oBufferId = funcAllocation->getBufferId(op); - int offset = -1; - if (oBufferId != Allocation::InvalidBufferId) - offset = funcAllocation->getOffset(oBufferId); - else if (op->getNumResults() == 1) { - Value value = op->getResult(0); - auto vBufferId = funcAllocation->getBufferId(value); - if (vBufferId != Allocation::InvalidBufferId) - offset = funcAllocation->getOffset(vBufferId); - } - if (offset == -1) - return; - op->setAttr("allocation.offset", - IntegerAttr::get(IntegerType::get(ctx, 32), offset)); - }); - return WalkResult::skip(); - }); - mod->setAttr("ttg.shared", - mlir::IntegerAttr::get(mlir::IntegerType::get(ctx, 32), - allocation.getSharedMemorySize())); -} - -} // namespace mlir::triton::gpu diff --git a/third_party/mthreads/lib/Conversion/TritonGPUToLLVM/CMakeLists.txt b/third_party/mthreads/lib/Conversion/TritonGPUToLLVM/CMakeLists.txt deleted file mode 100644 index 2c1d48b5f9..0000000000 --- a/third_party/mthreads/lib/Conversion/TritonGPUToLLVM/CMakeLists.txt +++ /dev/null @@ -1,41 +0,0 @@ -add_triton_library(TritonGPUToLLVM - DotOpToLLVM/FMA.cpp - DotOpToLLVM/FMADotUtility.cpp - AllocateSharedMemory.cpp - AllocateSharedMemoryUtility.cpp - AllocateWarpGroups.cpp - AssertOpToLLVM.cpp - ControlFlowOpToLLVM.cpp - ConvertLayoutOpToLLVM.cpp - ElementwiseOpToLLVM.cpp - FuncOpToLLVM.cpp - GatherOpToLLVM.cpp - GlobalScratchMemoryAllocation.cpp - HistogramOpToLLVM.cpp - MakeRangeOpToLLVM.cpp - MemoryOpToLLVM.cpp - PrintOpToLLVM.cpp - ReduceOpToLLVM.cpp - ScanOpToLLVM.cpp - SPMDOpToLLVM.cpp - TypeConverter.cpp - Utility.cpp - ViewOpToLLVM.cpp - WarpSpecializeUtility.cpp - - DEPENDS - TritonGPUConversionPassIncGen - - LINK_LIBS PUBLIC - MLIRIR - MLIRPass - MLIRGPUDialect - MLIRGPUToNVVMTransforms - MLIRGPUToROCDLTransforms - MLIRGPUTransforms - TritonAnalysis - TritonIR - TritonGPUIR - TritonGPUTransforms - TritonNvidiaGPUTransforms -) diff --git a/third_party/mthreads/lib/Conversion/TritonGPUToLLVM/ControlFlowOpToLLVM.cpp b/third_party/mthreads/lib/Conversion/TritonGPUToLLVM/ControlFlowOpToLLVM.cpp deleted file mode 100644 index 6e39bc067f..0000000000 --- a/third_party/mthreads/lib/Conversion/TritonGPUToLLVM/ControlFlowOpToLLVM.cpp +++ /dev/null @@ -1,171 +0,0 @@ -#include "triton/Conversion/TritonGPUToLLVM/PatternTritonGPUOpToLLVM.h" -#include "triton/Conversion/TritonGPUToLLVM/Utility.h" - -namespace { - -using namespace mlir; -using namespace mlir::triton; - -struct ReturnOpConversion : public ConvertOpToLLVMPattern { - using ConvertOpToLLVMPattern::ConvertOpToLLVMPattern; - - LogicalResult - matchAndRewrite(triton::ReturnOp op, OpAdaptor adaptor, - ConversionPatternRewriter &rewriter) const override { - auto funcOp = op->getParentOfType(); - auto loc = op.getLoc(); - auto b = TritonLLVMOpBuilder(loc, rewriter); - if (funcOp->hasAttr("nvvm.kernel")) { - // A GPU kernel - if (op.getNumOperands() > 0) { - return rewriter.notifyMatchFailure( - op, "Kernel functions do not support return with operands"); - } - rewriter.replaceOpWithNewOp(op, TypeRange(), ValueRange(), - op->getAttrs()); - } else { - // A device function - LLVM::ReturnOp newOp; - if (adaptor.getOperands().size() < 2) { - // Single or no return value. - newOp = LLVM::ReturnOp::create(rewriter, op.getLoc(), - adaptor.getOperands()); - } else { - // Pack the results into a struct. - auto packedResultsTy = this->getTypeConverter()->packFunctionResults( - funcOp.getResultTypes()); - Value packedResults = - LLVM::UndefOp::create(rewriter, op.getLoc(), packedResultsTy); - for (auto it : llvm::enumerate(adaptor.getOperands())) { - packedResults = b.insert_val(packedResultsTy, packedResults, - it.value(), it.index()); - } - newOp = LLVM::ReturnOp::create(rewriter, op.getLoc(), packedResults); - } - newOp->setAttrs(op->getAttrs()); - rewriter.replaceOp(op, newOp->getResults()); - } - return success(); - } -}; - -// CallOpInterfaceLowering is adapted from -// https://github.com/llvm/llvm-project/blob/fae656b2dd80246c3c6f01e9c77c49560368752c/mlir/lib/Conversion/FuncToLLVM/FuncToLLVM.cpp#L485 -struct CallOpConversion : public ConvertOpToLLVMPattern { - CallOpConversion(LLVMTypeConverter &converter, - const TargetInfoBase &targetInfo, PatternBenefit benefit) - : ConvertOpToLLVMPattern(converter, benefit), - targetInfo(targetInfo) {} - - LogicalResult - matchAndRewrite(triton::CallOp callOp, - typename triton::CallOp::Adaptor adaptor, - ConversionPatternRewriter &rewriter) const override { - auto promotedOperands = promoteOperands(callOp, adaptor, rewriter); - auto newCallOp = - convertCallOpToLLVMCallOp(callOp, promotedOperands, rewriter); - if (!newCallOp) - return failure(); - auto results = getCallOpResults(callOp, newCallOp, rewriter); - rewriter.replaceOp(callOp, results); - return success(); - } - -private: - SmallVector - promoteOperands(triton::CallOp callOp, - typename triton::CallOp::Adaptor adaptor, - ConversionPatternRewriter &rewriter) const { - // Get the last argument of the caller, which is the current stack pointer - // of shared memory and append it to the operands of the callOp. - auto loc = callOp.getLoc(); - auto b = TritonLLVMOpBuilder(loc, rewriter); - auto caller = callOp->getParentOfType(); - auto promotedOperands = this->getTypeConverter()->promoteOperands( - callOp.getLoc(), /*opOperands=*/callOp->getOperands(), - adaptor.getOperands(), rewriter); -#ifdef __TLE__ - // The call-site offset is sufficient to derive the callee frame base. - // Kernel callers do not carry `allocation.offset` themselves. - if (!callOp->hasAttr("allocation.offset")) { -#else - if (!caller->hasAttr("allocation.offset") || - !callOp->hasAttr("allocation.offset")) { -#endif - auto base = LLVM::getStackPointer(rewriter, caller); - promotedOperands.push_back(base); - } else { - auto base = LLVM::getSharedMemoryBase(loc, rewriter, targetInfo, callOp); - promotedOperands.push_back(base); - } - - auto opOffsetAttr = callOp->getAttrOfType( - "ttg.global_scratch_memory_offset"); - Value opOffsetVal; - if (opOffsetAttr) { - auto opOffset = opOffsetAttr.getValue().getZExtValue(); - opOffsetVal = b.i32_val(opOffset); - } - - promotedOperands.push_back(LLVM::getGlobalScratchPtr( - loc, rewriter, targetInfo, caller, opOffsetVal)); - promotedOperands.push_back( - LLVM::getProfileScratchPtr(loc, rewriter, caller)); - return promotedOperands; - } - - LLVM::CallOp - convertCallOpToLLVMCallOp(triton::CallOp callOp, - ArrayRef promotedOperands, - ConversionPatternRewriter &rewriter) const { - // Pack the result types into a struct. - Type packedResult = nullptr; - unsigned numResults = callOp.getNumResults(); - auto resultTypes = llvm::to_vector<4>(callOp.getResultTypes()); - - if (numResults != 0) { - if (!(packedResult = - this->getTypeConverter()->packFunctionResults(resultTypes))) - return nullptr; - } - auto newCallOp = LLVM::CallOp::create(rewriter, callOp.getLoc(), - packedResult ? TypeRange(packedResult) - : TypeRange(), - promotedOperands, callOp->getAttrs()); - newCallOp.getProperties().setOpBundleSizes( - rewriter.getDenseI32ArrayAttr({})); - newCallOp.getProperties().setOperandSegmentSizes( - {static_cast(promotedOperands.size()), 0}); - return newCallOp; - } - - SmallVector - getCallOpResults(triton::CallOp callOp, LLVM::CallOp newCallOp, - ConversionPatternRewriter &rewriter) const { - auto numResults = callOp.getNumResults(); - SmallVector results; - if (numResults < 2) { - // If < 2 results, packing did not do anything and we can just return. - results.append(newCallOp.result_begin(), newCallOp.result_end()); - } else { - // Otherwise, it had been converted to an operation producing a structure. - // Extract individual results from the structure and return them as list. - results.reserve(numResults); - for (unsigned i = 0; i < numResults; ++i) { - results.push_back(LLVM::ExtractValueOp::create( - rewriter, callOp.getLoc(), newCallOp->getResult(0), i)); - } - } - return results; - } - const TargetInfoBase &targetInfo; -}; - -} // namespace - -void mlir::triton::populateControlFlowOpToLLVMPattern( - LLVMTypeConverter &typeConverter, RewritePatternSet &patterns, - const TargetInfoBase &targetInfo, PatternBenefit benefit) { - patterns.add(typeConverter, benefit); - patterns.add(typeConverter, targetInfo, benefit); -} diff --git a/third_party/mthreads/lib/Conversion/TritonGPUToLLVM/DotOpToLLVM/FMADotUtility.cpp b/third_party/mthreads/lib/Conversion/TritonGPUToLLVM/DotOpToLLVM/FMADotUtility.cpp deleted file mode 100644 index fa2c814722..0000000000 --- a/third_party/mthreads/lib/Conversion/TritonGPUToLLVM/DotOpToLLVM/FMADotUtility.cpp +++ /dev/null @@ -1,170 +0,0 @@ -#include "triton/Conversion/TritonGPUToLLVM/FMADotUtility.h" -#include "triton/Conversion/TritonGPUToLLVM/Utility.h" - -using namespace mlir; - -namespace { - -/// OperandValueKey structure represents compile time part -/// of spatial coordinates of a value in a tensor. -/// -/// Every Value spatial coordinates(i.e. [batch;nonK;k]) in tensor can be -/// defined as: -/// -/// batch = (bRepIdx * CTABSize + bIdx) + (laneBCoord + warpBCoord) -/// nonK = (nonKRepIdx * CTANKSize + nonKIdx) + (laneNonKCoord + warpNonKCoord) -/// k = kIdx -/// -/// Where: -/// CTABSize, CTANKSize: constants; -/// laneBCoord, warpBCoord, laneNonKCoord, warpNonKCoord: runtime components; -/// bRepIdx, nonKRepIdx, bIdx, nonKIdx, kIdx: compile time components. -struct OperandValueKey { - unsigned bRepIdx, nonKRepIdx; - unsigned bIdx, nonKIdx, kIdx; - - bool operator==(const OperandValueKey &other) const { - return (bRepIdx == other.bRepIdx && nonKRepIdx == other.nonKRepIdx && - bIdx == other.bIdx && nonKIdx == other.nonKIdx && - kIdx == other.kIdx); - } -}; - -} // namespace - -template <> struct std::hash { - std::size_t operator()(const OperandValueKey &k) const { - return llvm::hash_combine(k.bRepIdx, k.nonKRepIdx, k.bIdx, k.nonKIdx, - k.kIdx); - } -}; - -namespace { - -using ValueTableFMA = std::unordered_map; - -ValueTableFMA getValueTableFromStructFMA( - Value val, ArrayRef perRepShape, ArrayRef repetitions, - unsigned kDim, unsigned nonKDim, ConversionPatternRewriter &rewriter, - Location loc, ArrayRef inRepOrder, ArrayRef repOrder) { - ValueTableFMA res; - auto elems = unpackLLElements(loc, val, rewriter); - assert(perRepShape.size() == 3); - auto numElemsRep = product(perRepShape); - assert(elems.size() == numElemsRep * product(repetitions)); - assert(kDim == 1 || kDim == 2); - assert(nonKDim == 1 || nonKDim == 2); - const unsigned bDim = 0; - - for (unsigned idx = 0; idx < elems.size(); ++idx) { - auto inRepLinearIdx = idx % numElemsRep; - auto repLinearIdx = idx / numElemsRep; - auto inRepSpatialIdx = - mlir::LLVM::delinearize(inRepLinearIdx, perRepShape, inRepOrder); - auto repSpatialIdx = - mlir::LLVM::delinearize(repLinearIdx, repetitions, repOrder); - OperandValueKey key{repSpatialIdx[0], repSpatialIdx[nonKDim], - inRepSpatialIdx[0], inRepSpatialIdx[nonKDim], - inRepSpatialIdx[kDim]}; - res[key] = elems[idx]; - } - return res; -} - -} // namespace - -namespace mlir::triton::gpu { - -LogicalResult parametricConvertFMADot(DotOp op, DotOp::Adaptor adaptor, - const LLVMTypeConverter *typeConverter, - ConversionPatternRewriter &rewriter, - FMAVectorMultiplier &multiplier) { - auto *ctx = rewriter.getContext(); - auto loc = op.getLoc(); - - auto A = op.getA(); - auto D = op.getResult(); - - auto aTensorTy = cast(A.getType()); - auto dTensorTy = cast(D.getType()); - - SmallVector aShapePerCTA = - expandMatrixShapeWithBatch(ArrayRef(getShapePerCTA(aTensorTy))); - auto dShapePerCTA = - expandMatrixShapeWithBatch(ArrayRef(getShapePerCTA(dTensorTy))); - - BlockedEncodingAttr dLayout = - cast(dTensorTy.getEncoding()); - // TODO process A and B operand separately - auto inRepOrder = expandMatrixOrderWithBatch(dLayout.getOrder()); - auto repOrder = expandMatrixOrderWithBatch(dLayout.getRepOrder()); - auto cc = unpackLLElements(loc, adaptor.getC(), rewriter); - - Value llA = adaptor.getA(); - Value llB = adaptor.getB(); - - auto sizePerThread = getContigPerThread(dTensorTy); - auto numElemsPerThread = product(sizePerThread); - SmallVector shapePerCTATile; - for (auto [reg, thread, warp] : - llvm::zip(sizePerThread, dLayout.getThreadsPerWarp(), - dLayout.getWarpsPerCTA())) { - shapePerCTATile.push_back(reg * thread * warp); - } - shapePerCTATile = expandMatrixShapeWithBatch(ArrayRef(shapePerCTATile)); - sizePerThread = expandMatrixShapeWithBatch(ArrayRef(sizePerThread)); - - unsigned K = aShapePerCTA[2]; - - unsigned threadTileShape[3]; - unsigned repetitions[3]; - for (int i = 0; i < 3; ++i) { - repetitions[i] = - ceil(dShapePerCTA[i], static_cast(shapePerCTATile[i])); - } - - auto has = getValueTableFromStructFMA( - llA, {sizePerThread[0], sizePerThread[1], K}, - {repetitions[0], repetitions[1], 1}, - /*kDim*/ 2, /*nonKDim*/ 1, rewriter, loc, inRepOrder, repOrder); - auto hbs = getValueTableFromStructFMA( - llB, {sizePerThread[0], K, sizePerThread[2]}, - {repetitions[0], 1, repetitions[2]}, - /*kDim*/ 1, /*nonKDim*/ 2, rewriter, loc, inRepOrder, repOrder); - - SmallVector acc = cc; - - for (unsigned bRep = 0; bRep < repetitions[0]; ++bRep) - for (unsigned mRep = 0; mRep < repetitions[1]; ++mRep) - for (unsigned nRep = 0; nRep < repetitions[2]; ++nRep) - for (unsigned b = 0; b < sizePerThread[0]; ++b) - for (unsigned m = 0; m < sizePerThread[1]; ++m) - for (unsigned n = 0; n < sizePerThread[2]; ++n) { - SmallVector multiDimAccumIdx = {b, m, n}; - unsigned linearInRepIdx = - LLVM::linearize(multiDimAccumIdx, sizePerThread, inRepOrder); - SmallVector multiDimRepIdx = {bRep, mRep, nRep}; - unsigned linearRepIdx = - LLVM::linearize(multiDimRepIdx, repetitions, repOrder); - unsigned linearAccumIdx = - linearInRepIdx + linearRepIdx * numElemsPerThread; - - SmallVector aOpVector; - SmallVector bOpVector; - - for (unsigned k = 0; k < K; ++k) { - aOpVector.push_back(has.at({bRep, mRep, b, m, k})); - bOpVector.push_back(hbs.at({bRep, nRep, b, n, k})); - } - - acc[linearAccumIdx] = multiplier.multiplyVectors( - aOpVector, bOpVector, acc[linearAccumIdx]); - } - - auto res = packLLElements(loc, typeConverter, acc, rewriter, dTensorTy); - rewriter.replaceOp(op, res); - - return success(); -} - -} // namespace mlir::triton::gpu diff --git a/third_party/mthreads/lib/Conversion/TritonGPUToLLVM/MakeRangeOpToLLVM.cpp b/third_party/mthreads/lib/Conversion/TritonGPUToLLVM/MakeRangeOpToLLVM.cpp deleted file mode 100644 index 8060b44312..0000000000 --- a/third_party/mthreads/lib/Conversion/TritonGPUToLLVM/MakeRangeOpToLLVM.cpp +++ /dev/null @@ -1,54 +0,0 @@ -#include "mlir/Conversion/LLVMCommon/Pattern.h" -#include "triton/Conversion/TritonGPUToLLVM/PatternTritonGPUOpToLLVM.h" -#include "triton/Conversion/TritonGPUToLLVM/TargetInfoBase.h" -#include "triton/Conversion/TritonGPUToLLVM/Utility.h" - -namespace { - -using namespace mlir; -using namespace mlir::triton; -struct MakeRangeOpConversion - : public ConvertOpToLLVMPattern { - MakeRangeOpConversion(LLVMTypeConverter &converter, - const TargetInfoBase &targetInfo, - PatternBenefit benefit) - : ConvertOpToLLVMPattern(converter, benefit), - targetInfo(targetInfo) {} - LogicalResult - matchAndRewrite(triton::MakeRangeOp op, OpAdaptor adaptor, - ConversionPatternRewriter &rewriter) const override { - Location loc = op->getLoc(); - auto b = TritonLLVMOpBuilder(loc, rewriter); - RankedTensorType ty = op.getType(); - auto shape = ty.getShape(); - auto layout = ty.getEncoding(); - auto elemTy = ty.getElementType(); - assert(elemTy.isInteger(32)); - Value start = createIndexAttrConstant(rewriter, loc, elemTy, op.getStart()); - auto idxs = emitIndices(loc, rewriter, targetInfo, layout, ty, true); - unsigned elems = idxs.size(); - SmallVector retVals(elems); - // TODO: slice layout has more elements than expected. - // Unexpected behavior for make range, but generally OK when followed by - // expand dims + broadcast. very weird behavior otherwise potentially. - for (const auto &multiDim : llvm::enumerate(idxs)) { - assert(multiDim.value().size() == 1); - retVals[multiDim.index()] = b.add(multiDim.value()[0], start); - } - auto typeConverter = getTypeConverter(); - Value result = packLLElements(loc, typeConverter, retVals, rewriter, ty); - rewriter.replaceOp(op, result); - return success(); - } - -private: - const TargetInfoBase &targetInfo; -}; - -} // namespace - -void mlir::triton::populateMakeRangeOpToLLVMPattern( - LLVMTypeConverter &typeConverter, const TargetInfoBase &targetInfo, - RewritePatternSet &patterns, PatternBenefit benefit) { - patterns.add(typeConverter, targetInfo, benefit); -} diff --git a/third_party/mthreads/lib/Conversion/TritonGPUToLLVM/PrintOpToLLVM.cpp b/third_party/mthreads/lib/Conversion/TritonGPUToLLVM/PrintOpToLLVM.cpp deleted file mode 100644 index e17b0e3ad3..0000000000 --- a/third_party/mthreads/lib/Conversion/TritonGPUToLLVM/PrintOpToLLVM.cpp +++ /dev/null @@ -1,243 +0,0 @@ -#include "mlir/Conversion/LLVMCommon/Pattern.h" -#include "mlir/Conversion/LLVMCommon/TypeConverter.h" -#include "mlir/IR/PatternMatch.h" -#include "triton/Conversion/TritonGPUToLLVM/PatternTritonGPUOpToLLVM.h" -#include "triton/Conversion/TritonGPUToLLVM/TargetInfoBase.h" -#include "triton/Conversion/TritonGPUToLLVM/Utility.h" -#include "triton/Dialect/Triton/IR/Dialect.h" - -namespace { - -// The input print op contains: -// - a "prefix" (string) specified by the user, and -// - one or more "operands" (tensors). -// -// For each operand, we print all of the values contained in this GPU thread, -// one per line, along with the index of the value in its tensor. -struct PrintOpConversion : public ConvertOpToLLVMPattern { - explicit PrintOpConversion(LLVMTypeConverter &typeConverter, - const TargetInfoBase &targetInfo, - PatternBenefit benefit) - : mlir::ConvertOpToLLVMPattern(typeConverter, benefit), - targetInfo(targetInfo) {} - - LogicalResult - matchAndRewrite(triton::PrintOp op, OpAdaptor adaptor, - ConversionPatternRewriter &rewriter) const override { - auto loc = op->getLoc(); - - std::array pid; - auto module = op->getParentOfType(); - for (auto axis : {ProgramIDDim::X, ProgramIDDim::Y, ProgramIDDim::Z}) - pid[(int)axis] = targetInfo.programId(rewriter, loc, module, axis); - - // Simple printf of a string without any tensors. - if (op.getNumOperands() == 0) { - std::string formatStr; - llvm::raw_string_ostream os(formatStr); - os << "pid (" << getFormatSubstr(pid[0]) << ", " - << getFormatSubstr(pid[1]) << ", " << getFormatSubstr(pid[2]) << ")" - << op.getPrefix(); - llPrintf(formatStr, {pid[0], pid[1], pid[2]}, {}, rewriter); - rewriter.eraseOp(op); - return success(); - } - - assert(op.getNumOperands() == op.getIsSigned().size()); - - for (size_t i = 0; i < op.getNumOperands(); i++) { - bool isSigned = op.getIsSigned()[i] > 0; - // Elements of the tensor that are resident in this GPU thread. - auto elems = unpackLLElements(loc, adaptor.getOperands()[i], rewriter); - - // Get the indices of `elems` within the tensor. Note that if `elems` - // has an "interesting" layout, then these will not be in any - // particularly nice order. - - // Extract the shape of the tensor being printed and use it to figure - // out how many digits we need for each of the dimensions. - SmallVector dimWidths; - SmallVector> indices; - if (auto rankedTy = - dyn_cast(op.getOperand(i).getType())) { - indices = emitIndices(loc, rewriter, targetInfo, rankedTy.getEncoding(), - rankedTy, true); - for (int64_t dim : rankedTy.getShape()) { - if (dim > 0) { - dimWidths.push_back(static_cast(std::ceil(std::log10(dim)))); - } else { - dimWidths.push_back(0); - } - } - } else { - // We're printing a scalar. - assert(elems.size() == 1); - indices.push_back({}); - } - - if (!elems.empty()) { - printTensor(op.getPrefix(), /*operand=*/i, - /*numOperands=*/op.getNumOperands(), elems, pid, indices, - dimWidths, op.getHex(), rewriter, isSigned); - } - } - rewriter.eraseOp(op); - return success(); - } - - void printTensor(StringRef prefixStr, size_t operand, size_t numOperands, - ArrayRef elems, std::array pid, - ArrayRef> indices, - ArrayRef dimWidths, bool hex, - ConversionPatternRewriter &rewriter, bool isSigned) const { - assert(!elems.empty()); - assert(elems.size() == indices.size()); - assert(dimWidths.size() == indices.front().size()); - - size_t rank = dimWidths.size(); - - // Format is: - // pid (, , ) idx (, , ...) (operand ) - // where we leave off "(operand )" if there's only one operand. - // - // The Python wrapper munges `prefix` so that it prints nicely (e.g. starts - // with " " and ends with ": "). - - Value formatStrValue; - int formatStrByteCount = 0; - for (int i = 0; i < elems.size(); i++) { - std::string formatStr; - llvm::raw_string_ostream os(formatStr); - - // nvptx printf can only accept 32 args; if we pass more than that, it - // will print garbage for the trailing args. - constexpr int kMaxPrintfOperands = 32; - SmallVector printfOperands; - - // TODO(jlebar): We really should pad the pid, but because the max pid is - // not known at compile-time, this would require nontrivial device-side - // work. - os << "pid ("; - for (int j = 0; j < pid.size(); j++) { - if (j != 0) { - os << ", "; - } - os << getFormatSubstr(pid[j]); - printfOperands.push_back(pid[j]); - } - os << ") "; - - // If `rank` is large enough, we could end up exceeding - // kMaxPrintfOperands. In that case, just truncate the index. - // (Subtract 2 because we're going to add two operands after the index.) - int maxAllowedRank = kMaxPrintfOperands - printfOperands.size() - 2; - - os << "idx ("; - const auto &index = indices[i]; - for (size_t dim = 0; dim < index.size(); dim++) { - if (dim != 0) { - os << ", "; - } - if (dim == maxAllowedRank) { - os << "... (truncated)"; - break; - } - os << getFormatSubstr(index[dim], /*hex=*/false, - /*width=*/dimWidths[dim]); - printfOperands.push_back(index[dim]); - } - os << ")" << prefixStr; - - if (numOperands > 1) { - os << "(operand " << operand << ") "; - } - - auto elem = elems[i]; - - os << getFormatSubstr(elem, hex, /*width=*/std::nullopt, isSigned); - printfOperands.push_back(elem); - - // It's the same format string each iteration, but it's a lot easier if we - // construct the format string at the same time as we populate - // printfOperands. But we don't want to create BLOCK_SIZE duplicate - // strings, so we cache the Value. - auto isSignedOperands = - llvm::SmallVector(printfOperands.size(), isSigned); - if (i == 0) { - formatStrValue = llPrintf(formatStr, printfOperands, isSignedOperands, - rewriter, &formatStrByteCount); - } else { - targetInfo.printf(rewriter, formatStrValue, formatStrByteCount, - printfOperands, isSignedOperands); - } - } - } - - std::string getFormatSubstr(Value value, bool hex = false, - std::optional width = std::nullopt, - bool isSigned = false) const { - Type type = value.getType(); - // If the `value` is a pointer, just return %p. - if (isa(type)) { - return "%p"; - } - // Hex is "0x%0nx" or "0x%0nllx", where n is the number of hex digits in the - // type (so 4 for fp16, 8 for int32, 16 for int64). - if (hex) { - // Ignore `width` for `hex` values, pad to typeWidth. - std::string ret = - "0x%0" + std::to_string(type.getIntOrFloatBitWidth() / 4); - if (type.getIntOrFloatBitWidth() > 32) { - ret += "ll"; - } - ret += "x"; - return ret; - } - - std::string prefix = "%"; - if (width.has_value()) { - prefix += std::to_string(*width); - } - - if (type.isBF16() || type.isF16() || type.isF32() || type.isF64()) { - return prefix + "f"; - } else if (type.isInteger()) { - if (type.getIntOrFloatBitWidth() == 64) - return prefix + (isSigned ? "lli" : "llu"); - else - return prefix + (isSigned ? "i" : "u"); - } - assert(false && "not supported type"); - return ""; - } - - // Returns a Value for the format string, which you can reuse. Writes the byte - // count for the string to |formatStrByteCount| if not null. - Value llPrintf(StringRef msg, ValueRange args, ArrayRef isSigned, - ConversionPatternRewriter &rewriter, - int *formatStrByteCount = nullptr) const { - assert(!msg.empty() && "printf with empty string not supported"); - llvm::SmallString<64> msgNewline(msg); - msgNewline.push_back('\n'); - msgNewline.push_back('\0'); - Value msgValue = - LLVM::addStringToModule(UnknownLoc::get(rewriter.getContext()), - rewriter, "printfFormat_", msgNewline); - targetInfo.printf(rewriter, msgValue, msgNewline.size_in_bytes(), args, - isSigned); - if (formatStrByteCount) - *formatStrByteCount = msgNewline.size_in_bytes(); - return msgValue; - } - -protected: - const TargetInfoBase &targetInfo; -}; - -} // namespace - -void mlir::triton::populatePrintOpToLLVMPattern( - LLVMTypeConverter &typeConverter, RewritePatternSet &patterns, - const TargetInfoBase &targetInfo, PatternBenefit benefit) { - patterns.add(typeConverter, targetInfo, benefit); -} diff --git a/third_party/mthreads/lib/Conversion/TritonGPUToLLVM/ReduceScanCommon.h b/third_party/mthreads/lib/Conversion/TritonGPUToLLVM/ReduceScanCommon.h deleted file mode 100644 index b132461761..0000000000 --- a/third_party/mthreads/lib/Conversion/TritonGPUToLLVM/ReduceScanCommon.h +++ /dev/null @@ -1,163 +0,0 @@ -#ifndef TRITON_CONVERSION_TRITONGPU_TO_LLVM_REDUCESCANCOMMON_H -#define TRITON_CONVERSION_TRITONGPU_TO_LLVM_REDUCESCANCOMMON_H - -// TODO: refactor so that it doesn't fail if Allocation.h -// is included after utility.h (due to conflict in `store` macro -// and -#include "mlir/Dialect/ControlFlow/IR/ControlFlowOps.h" -#include "mlir/Dialect/LLVMIR/LLVMDialect.h" -#include "mlir/Transforms/DialectConversion.h" - -// -#include "mlir/IR/TypeUtilities.h" -#include "triton/Conversion/TritonGPUToLLVM/Utility.h" -#include -#include - -#define DEBUG_TYPE "ttgpu_to_llvm" - -using namespace mlir; -using namespace mlir::triton; - -namespace mlir::triton { -class ReduceOp; -class ScanOp; - -inline SmallVector -inlineCombineBlock(ConversionPatternRewriter &rewriter, Block &combineBlock, - Block *insertionBlock, Block::iterator insertionPoint, - ValueRange combineArgs) { - auto returnOp = combineBlock.getTerminator(); - rewriter.inlineBlockBefore(&combineBlock, insertionBlock, insertionPoint, - combineArgs); - - auto results = SmallVector(returnOp->getOperands()); - - // Delete the terminator, which is no longer used - rewriter.eraseOp(returnOp); - return results; -} - -inline SmallVector applyCombineOp(Location loc, - ConversionPatternRewriter &rewriter, - Region &combineOp, ValueRange acc, - ValueRange cur, Value pred = {}) { - // Allows for passing an uninitialized acc and use cur as the neutral element - if (acc.size() == 0) { - return cur; - } - assert(cur.size() == acc.size()); - - // Create a new copy of the combine block, and try to speculatively inline it - Block *currentBlock = rewriter.getBlock(); - Region &parent = *currentBlock->getParent(); - - rewriter.cloneRegionBefore(combineOp, parent, - std::next(currentBlock->getIterator())); - Block &newCombine = *currentBlock->getNextNode(); - - llvm::SmallVector combineArgs(2 * acc.size()); - for (unsigned i = 0; i < acc.size(); ++i) { - combineArgs[i] = acc[i]; - combineArgs[acc.size() + i] = cur[i]; - } - - auto isRegionSpeculatable = - std::all_of(newCombine.begin(), newCombine.end(), - [](auto &op) { return isSpeculatable(&op); }); - - if (!pred || isRegionSpeculatable) { - // Fast path, region has no side effects so we can unconditionally execute - return inlineCombineBlock(rewriter, newCombine, currentBlock, - rewriter.getInsertionPoint(), combineArgs); - } - - // Slow case, create an if to only execute region when pred is true - // #currentBlock - // if (pred) { - // #newCombine - // results = combineOp(cur, acc) - // yield results - // } else { - // yield undef - // } - // #thenBlock - Block *thenBlock = - rewriter.splitBlock(currentBlock, rewriter.getInsertionPoint()); - - auto returnOp = newCombine.getTerminator(); - auto results = SmallVector(returnOp->getOperands()); - - rewriter.setInsertionPointToEnd(currentBlock); - SmallVector thenBlockArgs; - thenBlockArgs.reserve(results.size()); - for (auto result : results) { - auto ty = result.getType(); - auto undef = LLVM::UndefOp::create(rewriter, loc, ty); - thenBlockArgs.push_back(undef); - thenBlock->addArgument(ty, loc); - } - LLVM::CondBrOp::create(rewriter, loc, pred, &newCombine, combineArgs, - thenBlock, thenBlockArgs); - - // Split a block after the call. - rewriter.setInsertionPointToEnd(&newCombine); - rewriter.replaceOpWithNewOp(returnOp, results, thenBlock); - rewriter.setInsertionPointToStart(thenBlock); - return SmallVector(thenBlock->getArguments()); -} - -} // namespace mlir::triton - -template -class ConvertTritonGPUReduceScanToLLVMPattern - : public ConvertOpToLLVMPattern { -public: - // Make sure the class is only instantiated with Reduce and Scan - static_assert(std::is_same_v || - std::is_same_v); - - using ConvertOpToLLVMPattern::getTypeConverter; - using ConvertOpToLLVMPattern::ConvertOpToLLVMPattern; - - // Return the pointee type of the shared memory pointer for operand i. - Type getElementType(SourceOp op, int i) const { - auto ty = op.getInputTypes()[i].getElementType(); - return getTypeConverter()->convertType(ty); - } - - // Helper to compute the smem bases in both reductions and scans - SmallVector getSmemBases(SourceOp op, unsigned elems, - ConversionPatternRewriter &rewriter, - const TargetInfoBase &targetInfo) const { - auto loc = op.getLoc(); - auto b = TritonLLVMOpBuilder(loc, rewriter); - // indices will store the index of the op operands in descending order - // of their bitwidths - std::vector indices(op.getNumOperands()); - std::iota(indices.begin(), indices.end(), 0); - - std::sort(indices.begin(), indices.end(), [&](unsigned i, unsigned j) { - return op.getElementTypes()[i].getIntOrFloatBitWidth() > - op.getElementTypes()[j].getIntOrFloatBitWidth(); - }); - // Assign base index to each operand in their order in indices - std::map indexToBase; - auto basePtr = - LLVM::getSharedMemoryBase(loc, rewriter, targetInfo, op.getOperation()); - indexToBase[indices[0]] = basePtr; - for (unsigned i = 1; i < op.getNumOperands(); ++i) { - indexToBase[indices[i]] = - b.gep(basePtr.getType(), getElementType(op, indices[i - 1]), - indexToBase[indices[i - 1]], b.i32_val(elems)); - } - // smemBases[k] is the base pointer for the k-th operand - SmallVector smemBases(op.getNumOperands()); - for (unsigned i = 0; i < op.getNumOperands(); ++i) { - smemBases[i] = indexToBase[i]; - } - return smemBases; - } -}; - -#endif diff --git a/third_party/mthreads/lib/Conversion/TritonGPUToLLVM/SPMDOpToLLVM.cpp b/third_party/mthreads/lib/Conversion/TritonGPUToLLVM/SPMDOpToLLVM.cpp deleted file mode 100644 index 13b4f018f7..0000000000 --- a/third_party/mthreads/lib/Conversion/TritonGPUToLLVM/SPMDOpToLLVM.cpp +++ /dev/null @@ -1,37 +0,0 @@ -#include "triton/Conversion/TritonGPUToLLVM/PatternTritonGPUOpToLLVM.h" -#include "triton/Conversion/TritonGPUToLLVM/Utility.h" - -namespace { - -using namespace mlir; -using namespace mlir::triton; - -struct GetProgramIdOpConversion - : public ConvertOpToLLVMPattern { - explicit GetProgramIdOpConversion(LLVMTypeConverter &typeConverter, - const TargetInfoBase &targetInfo, - PatternBenefit benefit = 1) - : ConvertOpToLLVMPattern(typeConverter, benefit), - targetInfo(targetInfo) {} - - LogicalResult - matchAndRewrite(triton::GetProgramIdOp op, OpAdaptor adaptor, - ConversionPatternRewriter &rewriter) const override { - Value programId = targetInfo.programId( - rewriter, op->getLoc(), op->getParentOfType(), op.getAxis()); - rewriter.replaceOp(op, programId); - return success(); - } - -private: - const TargetInfoBase &targetInfo; -}; - -} // namespace - -void mlir::triton::populateSPMDOpToLLVMPattern(LLVMTypeConverter &typeConverter, - RewritePatternSet &patterns, - const TargetInfoBase &targetInfo, - PatternBenefit benefit) { - patterns.add(typeConverter, targetInfo, benefit); -} diff --git a/third_party/mthreads/lib/Conversion/TritonGPUToLLVM/TypeConverter.cpp b/third_party/mthreads/lib/Conversion/TritonGPUToLLVM/TypeConverter.cpp deleted file mode 100644 index f220ad3175..0000000000 --- a/third_party/mthreads/lib/Conversion/TritonGPUToLLVM/TypeConverter.cpp +++ /dev/null @@ -1,77 +0,0 @@ -#include "triton/Conversion/TritonGPUToLLVM/TypeConverter.h" - -#include "mlir/Support/LLVM.h" -#include "triton/Dialect/TritonNvidiaGPU/IR/Dialect.h" - -using namespace mlir; -using namespace mlir::triton; - -using ::mlir::triton::gpu::getTotalElemsPerThread; -using ::mlir::triton::gpu::MemDescType; - -TritonGPUToLLVMTypeConverter::TritonGPUToLLVMTypeConverter( - MLIRContext *ctx, const TargetInfoBase &targetInfo, - const DataLayoutAnalysis *analysis) - : TritonGPUToLLVMTypeConverter(ctx, LowerToLLVMOptions(ctx), targetInfo, - analysis) {} - -TritonGPUToLLVMTypeConverter::TritonGPUToLLVMTypeConverter( - MLIRContext *ctx, const LowerToLLVMOptions &options, - const TargetInfoBase &targetInfo, const DataLayoutAnalysis *analysis) - : LLVMTypeConverter(ctx, options, analysis) { - addConversion([ctx](triton::PointerType type) -> std::optional { - return LLVM::LLVMPointerType::get(ctx, type.getAddressSpace()); - }); - addConversion([ctx](TensorDescType type) -> std::optional { - return LLVM::LLVMPointerType::get(ctx, 0); - }); - addConversion([&](RankedTensorType type) -> std::optional { - return convertTritonTensorType(type, targetInfo); - }); - addConversion([&](MemDescType type) -> std::optional { - return convertMemDescType(type, targetInfo); - }); - addConversion([&](triton::gpu::AsyncTokenType type) -> std::optional { - return convertAsyncTokenType(type); - }); - - convertFP8Type(); -} - -Type TritonGPUToLLVMTypeConverter::convertTritonTensorType( - RankedTensorType type, const TargetInfoBase &targetInfo) { - auto ctx = type.getContext(); - Type eltType = convertType(type.getElementType()); - unsigned numElementsPerThread = getTotalElemsPerThread(type); - SmallVector types(numElementsPerThread, eltType); - return LLVM::LLVMStructType::getLiteral(ctx, types); -} - -Type TritonGPUToLLVMTypeConverter::convertMemDescType( - MemDescType type, const TargetInfoBase &targetInfo) { - auto ctx = type.getContext(); - // base ptr - auto ptrType = LLVM::LLVMPointerType::get( - ctx, targetInfo.getAddressSpace(type.getMemorySpace())); - - if (isa( - type.getEncoding())) { - return ptrType; - } - - SmallVector types; - types.push_back(ptrType); - auto rank = type.getRank(); - // offsets - for (auto i = 0; i < rank; i++) { - types.push_back(IntegerType::get(ctx, 32)); - } - return LLVM::LLVMStructType::getLiteral(ctx, types); -} - -Type TritonGPUToLLVMTypeConverter::convertAsyncTokenType( - triton::gpu::AsyncTokenType type) { - return IntegerType::get(type.getContext(), 32); -} diff --git a/third_party/mthreads/lib/Conversion/TritonInstrumentToLLVM/CMakeLists.txt b/third_party/mthreads/lib/Conversion/TritonInstrumentToLLVM/CMakeLists.txt deleted file mode 100644 index 5a3c379304..0000000000 --- a/third_party/mthreads/lib/Conversion/TritonInstrumentToLLVM/CMakeLists.txt +++ /dev/null @@ -1,12 +0,0 @@ -add_triton_library(TritonInstrumentToLLVM - InstrumentationToLLVM.cpp - - LINK_LIBS PUBLIC - MLIRIR - MLIRPass - TritonIR - TritonGPUIR - TritonInstrumentIR - TritonNvidiaGPUIR - NVGPUIR -) diff --git a/third_party/mthreads/lib/Conversion/TritonToTritonGPU/CMakeLists.txt b/third_party/mthreads/lib/Conversion/TritonToTritonGPU/CMakeLists.txt deleted file mode 100644 index de600e21b5..0000000000 --- a/third_party/mthreads/lib/Conversion/TritonToTritonGPU/CMakeLists.txt +++ /dev/null @@ -1,26 +0,0 @@ -if(FLAGTREE_MTHREADS_TLE) - set(_MUSATLE_DEPS MUSATLETableGen) - set(_MUSATLE_LIBS MUSATLEIR) -else() - set(_MUSATLE_DEPS "") - set(_MUSATLE_LIBS "") -endif() - -add_triton_library(TritonToTritonGPU - RelayoutTritonGPU.cpp - TritonGPUConversion.cpp - TritonToTritonGPUPass.cpp - - DEPENDS - TritonConversionPassIncGen - ${_MUSATLE_DEPS} - - LINK_LIBS PUBLIC - MLIRIR - MLIRPass - MLIRTransforms - TritonIR - ProtonIR - TritonGPUIR - ${_MUSATLE_LIBS} -) diff --git a/third_party/mthreads/lib/Dialect/CMakeLists.txt b/third_party/mthreads/lib/Dialect/CMakeLists.txt deleted file mode 100644 index 19ca22ec3b..0000000000 --- a/third_party/mthreads/lib/Dialect/CMakeLists.txt +++ /dev/null @@ -1,7 +0,0 @@ -add_subdirectory(Triton) -add_subdirectory(TritonGPU) -add_subdirectory(TritonNvidiaGPU) -add_subdirectory(TritonInstrument) -add_subdirectory(Gluon) -add_subdirectory(NVGPU) -add_subdirectory(NVWS) diff --git a/third_party/mthreads/lib/Dialect/Gluon/CMakeLists.txt b/third_party/mthreads/lib/Dialect/Gluon/CMakeLists.txt deleted file mode 100644 index 9f57627c32..0000000000 --- a/third_party/mthreads/lib/Dialect/Gluon/CMakeLists.txt +++ /dev/null @@ -1,2 +0,0 @@ -add_subdirectory(IR) -add_subdirectory(Transforms) diff --git a/third_party/mthreads/lib/Dialect/Gluon/IR/CMakeLists.txt b/third_party/mthreads/lib/Dialect/Gluon/IR/CMakeLists.txt deleted file mode 100644 index 315f033e22..0000000000 --- a/third_party/mthreads/lib/Dialect/Gluon/IR/CMakeLists.txt +++ /dev/null @@ -1,10 +0,0 @@ -add_triton_library(GluonIR - Dialect.cpp - - DEPENDS - GluonTableGen - - LINK_LIBS PUBLIC - TritonIR - TritonGPUIR -) diff --git a/third_party/mthreads/lib/Dialect/Gluon/IR/Dialect.cpp b/third_party/mthreads/lib/Dialect/Gluon/IR/Dialect.cpp deleted file mode 100644 index 0a18ec8522..0000000000 --- a/third_party/mthreads/lib/Dialect/Gluon/IR/Dialect.cpp +++ /dev/null @@ -1,138 +0,0 @@ -#include "triton/Dialect/Gluon/IR/Dialect.h" - -#include "mlir/Support/LLVM.h" -#include "triton/Dialect/Triton/IR/Interfaces.h" -#include "llvm/ADT/TypeSwitch.h" - -using namespace mlir; -using namespace mlir::triton::gpu; -namespace gluon = mlir::triton::gluon; - -#define GET_ATTRDEF_CLASSES -#include "triton/Dialect/Gluon/IR/Dialect.cpp.inc" -#include "triton/Dialect/Gluon/IR/GluonAttrDefs.cpp.inc" - -#define GET_OP_CLASSES -#include "triton/Dialect/Gluon/IR/Ops.cpp.inc" - -namespace { - -// Layout inference for AutoEncodingAttr -> always propagate AutoEncodingAttr to -// results -struct GluonInferLayoutInterface : public triton::DialectInferLayoutInterface { - using DialectInferLayoutInterface::DialectInferLayoutInterface; - - LogicalResult inferAutoEncoding(Attribute operandEncoding, - Attribute &resultEncoding) const { - if (!isa( - operandEncoding)) - return failure(); - resultEncoding = operandEncoding; - return success(); - } - - LogicalResult - inferReduceOpEncoding(Attribute operandEncoding, unsigned axis, - Attribute &resultEncoding, - std::optional loc) const override { - return inferAutoEncoding(operandEncoding, resultEncoding); - } - - LogicalResult - inferTransOpEncoding(Attribute operandEncoding, ArrayRef shape, - ArrayRef order, Attribute &resultEncoding, - std::optional loc) const override { - return inferAutoEncoding(operandEncoding, resultEncoding); - } - - LogicalResult - inferExpandDimsOpEncoding(Attribute operandEncoding, unsigned axis, - Attribute &resultEncoding, - std::optional location) const override { - return inferAutoEncoding(operandEncoding, resultEncoding); - } - - LogicalResult - inferDotOpEncoding(Attribute operandEncoding, unsigned opIdx, - Attribute resultEncoding, - std::optional location) const override { - return inferAutoEncoding(operandEncoding, resultEncoding); - } - - LogicalResult - verifyDotOpEncodingCompatibility(Operation *op, Attribute operandEncodingA, - Attribute operandEncodingB) const override { - return success(); - } - - LogicalResult - verifyLayoutsAreEqual(ArrayRef shape, Attribute expected, - Attribute got, - std::optional loc) const override { - return success(expected == got); - } - - LogicalResult - inferReshapeOpEncoding(ArrayRef srcShape, Attribute srcEnc, - ArrayRef dstShape, Attribute &dstEnc, - std::optional loc) const override { - return inferAutoEncoding(srcEnc, dstEnc); - } - - LogicalResult - inferDefaultJoinOpEncoding(Attribute srcEnc, Attribute &dstEnc, - ArrayRef shape, - std::optional loc) const override { - return inferAutoEncoding(srcEnc, dstEnc); - } - - LogicalResult - inferSplitOpEncoding(Attribute srcEnc, Attribute &dstEnc, - ArrayRef shape, - std::optional loc) const override { - return inferAutoEncoding(srcEnc, dstEnc); - } - - LogicalResult - inferFp4ToFpOpEncoding(ArrayRef shape, int axis, Attribute srcEnc, - Attribute &dstEnc, bool fwdInference, - std::optional loc) const override { - return inferAutoEncoding(srcEnc, dstEnc); - } -}; -} // namespace - -namespace mlir::triton::gluon { - -void GluonDialect::initialize() { - addAttributes< -#define GET_ATTRDEF_LIST -#include "triton/Dialect/Gluon/IR/GluonAttrDefs.cpp.inc" - >(); - addOperations< -#define GET_OP_LIST -#include "triton/Dialect/Gluon/IR/Ops.cpp.inc" - >(); - addInterfaces(); - addInterfaces(); -} - -void SetAutoLayoutOp::build(OpBuilder &builder, OperationState &state, - Attribute enc, Value value) { - auto resTy = cast(value.getType()).cloneWithEncoding(enc); - return build(builder, state, resTy, value); -} - -LogicalResult SetAutoLayoutOp::verify() { - if (!isa(getSrc().getType().getEncoding())) { - return emitOpError("input tensor must have an auto layout type"); - } - auto dstEncoding = getType().getEncoding(); - if (!dstEncoding) - return emitOpError("result tensor must have an encoding"); - if (isa(dstEncoding)) - return emitOpError("result type must not be auto layout"); - return success(); -} - -} // namespace mlir::triton::gluon diff --git a/third_party/mthreads/lib/Dialect/Gluon/Transforms/CMakeLists.txt b/third_party/mthreads/lib/Dialect/Gluon/Transforms/CMakeLists.txt deleted file mode 100644 index 0e43d594c2..0000000000 --- a/third_party/mthreads/lib/Dialect/Gluon/Transforms/CMakeLists.txt +++ /dev/null @@ -1,17 +0,0 @@ -add_triton_library(GluonTransforms - Canonicalize.cpp - Inline.cpp - ResolveAutoEncodings.cpp - SimplifyControlFlow.cpp - InferCoalescedEncodings.cpp - InferLayoutUtils.cpp - - DEPENDS - GluonTransformsIncGen - - LINK_LIBS PUBLIC - TritonIR - TritonGPUIR - GluonIR - MLIRTransformUtils -) diff --git a/third_party/mthreads/lib/Dialect/Gluon/Transforms/InferLayoutUtils.cpp b/third_party/mthreads/lib/Dialect/Gluon/Transforms/InferLayoutUtils.cpp deleted file mode 100644 index bff4e64a4b..0000000000 --- a/third_party/mthreads/lib/Dialect/Gluon/Transforms/InferLayoutUtils.cpp +++ /dev/null @@ -1,251 +0,0 @@ -#include "triton/Dialect/Gluon/Transforms/InferLayoutUtils.h" -#include "mlir/IR/Attributes.h" -#include "mlir/IR/BuiltinAttributes.h" -#include "mlir/IR/Visitors.h" -#include "mlir/Support/LLVM.h" -#include "triton/Dialect/Gluon/IR/Dialect.h" -#include "triton/Dialect/Gluon/Transforms/Passes.h" -#include "triton/Dialect/TritonGPU/IR/Dialect.h" -#include "triton/Dialect/TritonGPU/Transforms/Utility.h" -#include "llvm/ADT/MapVector.h" -#include "llvm/ADT/PriorityWorklist.h" -#include "llvm/Support/Debug.h" -#include "llvm/Support/LogicalResult.h" -#include "llvm/Support/raw_ostream.h" -#include "llvm/Support/xxhash.h" - -#define DEBUG_TYPE "gluon-infer-layout-utils" -#define DBGS() (llvm::dbgs() << "[" DEBUG_TYPE "]: ") -#define LDBG(X) LLVM_DEBUG(DBGS() << X << "\n") - -namespace mlir::triton::gluon { - -namespace { -struct LayoutInfo { - Attribute encoding; - // Some operations can infer one of many encodings, - // we model this by setting the mayVary flag on encodings - // derived from these ops. - // If "may vary" is set then we allow conflicts, and when - // resolving conflicts we prefer encodings that are not allowed to vary. - bool mayVary = false; - - operator bool() { return bool(encoding); } -}; - -uint64_t hashWithMemo(Attribute attr, - llvm::MapVector &hashMemo) { - auto it = hashMemo.find(attr); - if (it != hashMemo.end()) { - return it->second; - } - - // llvm::hash_value is not stable, so instead we hash the string repr of the - // attribute - std::string str; - llvm::raw_string_ostream os(str); - attr.print(os); - auto hash = llvm::xxh3_64bits(str); - hashMemo.try_emplace(attr, hash); - return hash; -} - -bool compare(Attribute a, Attribute b, - llvm::MapVector &hashMemo) { - if (a == b) - return false; - - return hashWithMemo(a, hashMemo) > hashWithMemo(b, hashMemo); -} - -LayoutInfo combineInfo(LayoutInfo lhs, LayoutInfo rhs, Operation *op, - llvm::MapVector &hashMemo) { - // Sort inputs so this operation is commutative - if (compare(lhs.encoding, rhs.encoding, hashMemo)) { - std::swap(lhs, rhs); - } - if (lhs.mayVary) - return rhs; - if (rhs.mayVary) - return lhs; - if (lhs.encoding == rhs.encoding) - return lhs; - op->emitOpError("found conflicting encodings for value:\n ") - << lhs.encoding << "\nand\n " << rhs.encoding; - return {}; -} - -bool encodingsMayVary(Operation *op) { - return isa(op); -} - -LogicalResult -updateEncoding(ArrayRef values, LayoutInfo info, FuncOp *func, - llvm::MapVector &valueToEncoding, - llvm::PriorityWorklist &worklist, - llvm::MapVector &hashMemo) { - for (auto value : values) { - auto [it, inserted] = valueToEncoding.insert({value, info}); - if (!inserted) { - auto defOp = value.getDefiningOp(); - auto op = defOp ? defOp : func->getOperation(); - auto combine = combineInfo(it->second, info, op, hashMemo); - if (!combine) - return failure(); - if (combine == it->second) - continue; - it->second = combine; - } - LLVM_DEBUG({ - DBGS() << "Setting value:\n\t" << value << "\nto encoding:\n\t" - << it->second.encoding << "\n"; - }); - worklist.insert(value); - } - return success(); -} -} // namespace - -LogicalResult inferLayout( - FuncOp func, llvm::function_ref typeCheck, - const llvm::SmallVector> &seedEncodings) { - // Disallow auto encoding accross function call boundaries - for (auto argTy : func.getArgumentTypes()) { - if (typeCheck(argTy)) { - return func->emitError( - "Functions taking auto encoding must be fully inlined"); - } - } - for (auto resultTy : func.getResultTypes()) { - if (typeCheck(resultTy)) - return func->emitError( - "Functions returning auto encoding must be fully inlined"); - } - - // set seed - llvm::MapVector valueToEncoding; - llvm::PriorityWorklist worklist; - llvm::MapVector hashMemo; - for (auto &[value, encoding] : seedEncodings) { - if (failed(updateEncoding({value}, LayoutInfo{encoding, false}, &func, - valueToEncoding, worklist, hashMemo))) - return failure(); - } - - // Propagate encodings through the graph until fixed point, or conflict - while (!worklist.empty()) { - auto val = worklist.pop_back_val(); - auto info = valueToEncoding[val]; - assert(info); - - // Propagate to users - for (OpOperand &use : val.getUses()) { - auto op = use.getOwner(); - if (isa(op)) { - auto offset = 3 * isa(op); - auto tiedArgs = getTiedArgs(op, use.getOperandNumber() - offset); - if (failed(updateEncoding(tiedArgs, info, &func, valueToEncoding, - worklist, hashMemo))) - return failure(); - } else if (isa(op)) { - auto tiedArgs = getTiedArgs(op, use.getOperandNumber()); - if (failed(updateEncoding(tiedArgs, info, &func, valueToEncoding, - worklist, hashMemo))) - return failure(); - } else { - auto dstEnc = inferDstEncoding(op, info.encoding); - if (dstEnc) { - bool mayVary = info.mayVary || encodingsMayVary(op); - LayoutInfo dstInfo{dstEnc, mayVary}; - if (failed(updateEncoding(llvm::to_vector_of(op->getResults()), - dstInfo, &func, valueToEncoding, worklist, - hashMemo))) - return failure(); - } - } - } - - // Propagate to defining ops - if (auto opResult = dyn_cast(val)) { - auto definingOp = opResult.getOwner(); - if (isa(definingOp)) { - auto tiedArgs = getTiedArgs(definingOp, opResult.getResultNumber()); - if (failed(updateEncoding(tiedArgs, info, &func, valueToEncoding, - worklist, hashMemo))) - return failure(); - } else { - auto srcEncoding = inferSrcEncoding(definingOp, info.encoding); - if (srcEncoding) { - bool mayVary = info.mayVary || encodingsMayVary(definingOp); - LayoutInfo srcInfo{srcEncoding, mayVary}; - llvm::SmallVector tensorOperands; - for (auto operand : definingOp->getOperands()) - if (isa(operand.getType())) - tensorOperands.push_back(operand); - - if (failed(updateEncoding(tensorOperands, srcInfo, &func, - valueToEncoding, worklist, hashMemo))) - return failure(); - } - } - } else if (auto blockArg = dyn_cast(val)) { - auto parentOp = blockArg.getOwner()->getParentOp(); - if (isa(parentOp)) { - auto offset = isa(parentOp); - auto tiedArgs = getTiedArgs(parentOp, blockArg.getArgNumber() - offset); - if (failed(updateEncoding(tiedArgs, info, &func, valueToEncoding, - worklist, hashMemo))) - return failure(); - } - } - } - - // Transfer propagated encodings into the graph - auto ctx = func.getContext(); - for (auto &[val, info] : valueToEncoding) { - assert(typeCheck(val.getType())); - auto existingTy = cast(val.getType()); - auto ty = existingTy.cloneWithEncoding(info.encoding); - val.setType(ty); - - if (auto opResult = dyn_cast(val)) { - if (auto constantOp = dyn_cast(opResult.getOwner())) { - auto value = cast(constantOp.getValueAttr()); - auto newValue = - SplatElementsAttr::get(ty, value.getSplatValue()); - constantOp.setValueAttr(newValue); - } - } - } - return success(); -} - -LogicalResult doubleCheckEncodings(ModuleOp &mod, - llvm::function_ref typeCheck) { - auto res = mod.walk([&](Operation *op) -> WalkResult { - for (auto resTy : op->getResultTypes()) { - if (typeCheck(resTy)) { - return op->emitOpError("Failed to infer return type"); - } - } - return success(); - }); - if (res.wasInterrupted()) - return failure(); - - res = mod.walk([&](Block *block) -> WalkResult { - for (auto argTy : block->getArgumentTypes()) { - if (typeCheck(argTy)) { - return block->getParentOp()->emitError( - "Failed to infer block argument type"); - } - } - return success(); - }); - if (res.wasInterrupted()) - return failure(); - return success(); -} - -} // namespace mlir::triton::gluon diff --git a/third_party/mthreads/lib/Dialect/Gluon/Transforms/Inline.cpp b/third_party/mthreads/lib/Dialect/Gluon/Transforms/Inline.cpp deleted file mode 100644 index 0dd7d26c73..0000000000 --- a/third_party/mthreads/lib/Dialect/Gluon/Transforms/Inline.cpp +++ /dev/null @@ -1,29 +0,0 @@ -#include "triton/Dialect/Gluon/Transforms/Passes.h" - -#include "mlir/Pass/Pass.h" -#include "mlir/Pass/PassManager.h" -#include "mlir/Transforms/Passes.h" - -using namespace mlir; -using namespace triton; -namespace gluon = mlir::triton::gluon; - -namespace mlir::triton::gluon { -#define GEN_PASS_DEF_GLUONINLINE -#include "triton/Dialect/Gluon/Transforms/Passes.h.inc" -} // namespace mlir::triton::gluon - -namespace { -struct Inline : public gluon::impl::GluonInlineBase { - void runOnOperation() override; -}; -} // namespace - -void Inline::runOnOperation() { - mlir::PassManager pm(&getContext()); - pm.addPass(createInlinerPass(/*opPipelines=*/{}, [](OpPassManager &pm) { - pm.addPass(gluon::createGluonSimplifyControlFlow()); - })); - if (failed(pm.run(getOperation()))) - return signalPassFailure(); -} diff --git a/third_party/mthreads/lib/Dialect/Gluon/Transforms/ResolveAutoEncodings.cpp b/third_party/mthreads/lib/Dialect/Gluon/Transforms/ResolveAutoEncodings.cpp deleted file mode 100644 index c7b775cb7a..0000000000 --- a/third_party/mthreads/lib/Dialect/Gluon/Transforms/ResolveAutoEncodings.cpp +++ /dev/null @@ -1,71 +0,0 @@ -#include "triton/Dialect/Gluon/IR/Dialect.h" -#include "triton/Dialect/Gluon/Transforms/InferLayoutUtils.h" -#include "triton/Dialect/Gluon/Transforms/Passes.h" -#include "llvm/ADT/MapVector.h" -#include "llvm/ADT/PriorityWorklist.h" -#include "llvm/Support/Debug.h" -#include "llvm/Support/LogicalResult.h" - -namespace ttg = mlir::triton::gpu; - -namespace mlir::triton::gluon { - -#define GEN_PASS_DEF_GLUONRESOLVEAUTOENCODINGSPASS -#include "triton/Dialect/Gluon/Transforms/Passes.h.inc" - -#define DEBUG_TYPE "gluon-resolve-auto-encodings" -#define DBGS() (llvm::dbgs() << "[" DEBUG_TYPE "]: ") -#define LDBG(X) LLVM_DEBUG(DBGS() << X << "\n") - -namespace { -bool isAutoEncodingTensorType(Type ty) { - auto tensorTy = dyn_cast(ty); - return tensorTy && isa(tensorTy.getEncoding()); -} -LogicalResult inferAutoLayout(ModuleOp &mod) { - for (auto &op : *mod.getBody()) { - auto func = dyn_cast(&op); - if (!func) - continue; - - // Set seed values from set_auto_layout ops - llvm::SmallVector> seedEncodings; - func.walk([&](gluon::SetAutoLayoutOp op) { - seedEncodings.push_back({op.getSrc(), op.getType().getEncoding()}); - }); - - if (failed(inferLayout(func, isAutoEncodingTensorType, seedEncodings))) - return failure(); - } - return success(); -} -} // anonymous namespace - -class GluonResolveAutoEncodingsPass - : public impl::GluonResolveAutoEncodingsPassBase< - GluonResolveAutoEncodingsPass> { -public: - using BaseT = - impl::GluonResolveAutoEncodingsPassBase; - using BaseT::BaseT; - - void runOnOperation() override { - MLIRContext *context = &getContext(); - ModuleOp m = getOperation(); - - // Do layout inference - if (failed(inferAutoLayout(m))) - return signalPassFailure(); - - // Cleanup set_auto_layout ops - m.walk([&](gluon::SetAutoLayoutOp op) { - assert(op.getSrc().getType() == op.getType()); - op.getResult().replaceAllUsesWith(op.getSrc()); - op->erase(); - }); - - if (failed(doubleCheckEncodings(m, isAutoEncodingTensorType))) - return signalPassFailure(); - } -}; -} // namespace mlir::triton::gluon diff --git a/third_party/mthreads/lib/Dialect/Gluon/Transforms/SimplifyControlFlow.cpp b/third_party/mthreads/lib/Dialect/Gluon/Transforms/SimplifyControlFlow.cpp deleted file mode 100644 index c0a6b40f68..0000000000 --- a/third_party/mthreads/lib/Dialect/Gluon/Transforms/SimplifyControlFlow.cpp +++ /dev/null @@ -1,49 +0,0 @@ -#include "mlir/IR/OperationSupport.h" -#include "triton/Dialect/Gluon/Transforms/Passes.h" - -#include "triton/Dialect/TritonGPU/Transforms/Utility.h" - -#include "mlir/Dialect/Arith/IR/Arith.h" -#include "mlir/Dialect/ControlFlow/IR/ControlFlow.h" -#include "mlir/Dialect/SCF/IR/SCF.h" -#include "mlir/Pass/Pass.h" -#include "mlir/Transforms/GreedyPatternRewriteDriver.h" - -using namespace mlir; -using namespace triton; - -namespace mlir::triton::gluon { -#define GEN_PASS_DEF_GLUONSIMPLIFYCONTROLFLOW -#include "triton/Dialect/Gluon/Transforms/Passes.h.inc" -} // namespace mlir::triton::gluon - -namespace { -struct SimplifyControlFlow - : public gluon::impl::GluonSimplifyControlFlowBase { - void runOnOperation() override; -}; -} // namespace - -void SimplifyControlFlow::runOnOperation() { - MLIRContext *ctx = &getContext(); - RewritePatternSet patterns(&getContext()); - - // Populate `scf` and `cf` canonicalizers. - ctx->getLoadedDialect()->getCanonicalizationPatterns( - patterns); - ctx->getLoadedDialect()->getCanonicalizationPatterns( - patterns); - for (mlir::RegisteredOperationName op : ctx->getRegisteredOperationsByDialect( - scf::SCFDialect::getDialectNamespace())) - op.getCanonicalizationPatterns(patterns, ctx); - for (mlir::RegisteredOperationName op : ctx->getRegisteredOperationsByDialect( - cf::ControlFlowDialect::getDialectNamespace())) - op.getCanonicalizationPatterns(patterns, ctx); - populateForOpDeadArgumentElimination(patterns); - - GreedyRewriteConfig config; - // This is intended to run before AutoLayouts are resolved, in which case - // CSEing constants can lead to additional layout conflicts. - config.enableConstantCSE(false); - (void)applyPatternsGreedily(getOperation(), std::move(patterns), config); -} diff --git a/third_party/mthreads/lib/Dialect/NVGPU/CMakeLists.txt b/third_party/mthreads/lib/Dialect/NVGPU/CMakeLists.txt deleted file mode 100644 index f33061b2d8..0000000000 --- a/third_party/mthreads/lib/Dialect/NVGPU/CMakeLists.txt +++ /dev/null @@ -1 +0,0 @@ -add_subdirectory(IR) diff --git a/third_party/mthreads/lib/Dialect/NVWS/CMakeLists.txt b/third_party/mthreads/lib/Dialect/NVWS/CMakeLists.txt deleted file mode 100644 index 9f57627c32..0000000000 --- a/third_party/mthreads/lib/Dialect/NVWS/CMakeLists.txt +++ /dev/null @@ -1,2 +0,0 @@ -add_subdirectory(IR) -add_subdirectory(Transforms) diff --git a/third_party/mthreads/lib/Dialect/Triton/CMakeLists.txt b/third_party/mthreads/lib/Dialect/Triton/CMakeLists.txt deleted file mode 100644 index 9f57627c32..0000000000 --- a/third_party/mthreads/lib/Dialect/Triton/CMakeLists.txt +++ /dev/null @@ -1,2 +0,0 @@ -add_subdirectory(IR) -add_subdirectory(Transforms) diff --git a/third_party/mthreads/lib/Dialect/Triton/IR/CMakeLists.txt b/third_party/mthreads/lib/Dialect/Triton/IR/CMakeLists.txt deleted file mode 100644 index 1662b94968..0000000000 --- a/third_party/mthreads/lib/Dialect/Triton/IR/CMakeLists.txt +++ /dev/null @@ -1,23 +0,0 @@ -set(LLVM_TARGET_DEFINITIONS Canonicalize.td) -mlir_tablegen(TritonCanonicalize.inc -gen-rewriters) -add_public_tablegen_target(TritonCanonicalizeIncGen) - -add_triton_library(TritonIR - Dialect.cpp - DiscardableAttributes.cpp - Ops.cpp - Traits.cpp - Types.cpp - OpInterfaces.cpp - Utility.cpp - - DEPENDS - TritonTableGen - TritonCanonicalizeIncGen - - LINK_LIBS PUBLIC - MLIRIR - MLIRArithDialect - MLIRMathDialect - MLIRSCFDialect -) diff --git a/third_party/mthreads/lib/Dialect/Triton/IR/Canonicalize.td b/third_party/mthreads/lib/Dialect/Triton/IR/Canonicalize.td deleted file mode 100644 index dc37710333..0000000000 --- a/third_party/mthreads/lib/Dialect/Triton/IR/Canonicalize.td +++ /dev/null @@ -1,17 +0,0 @@ -#ifndef TT_PATTERNS -#define TT_PATTERNS - -include "mlir/IR/PatternBase.td" -include "triton/Dialect/Triton/IR/TritonOps.td" - -// broadcast(splat(x)) -> splat(x) -def BroadcastSplatPattern : - Pat<(TT_BroadcastOp (TT_SplatOp $x)), - (TT_SplatOp $x)>; - -// broadcast(broadcast(x)) -> broadcast(x) -def BroadcastBroadcastPattern : - Pat<(TT_BroadcastOp (TT_BroadcastOp $x)), - (TT_BroadcastOp $x)>; - -#endif diff --git a/third_party/mthreads/lib/Dialect/Triton/IR/Dialect.cpp b/third_party/mthreads/lib/Dialect/Triton/IR/Dialect.cpp deleted file mode 100644 index 9073f423f9..0000000000 --- a/third_party/mthreads/lib/Dialect/Triton/IR/Dialect.cpp +++ /dev/null @@ -1,77 +0,0 @@ -#include "triton/Dialect/Triton/IR/Dialect.h" -#include "triton/Dialect/Triton/IR/Interfaces.h" -#include "triton/Dialect/Triton/IR/Types.h" - -#include "mlir/Dialect/ControlFlow/IR/ControlFlowOps.h" -#include "mlir/Dialect/UB/IR/UBOps.h" -#include "llvm/ADT/StringSwitch.h" -#include "llvm/ADT/TypeSwitch.h" - -#include "triton/Dialect/Triton/IR/AttrInterfaces.cpp.inc" -#include "triton/Dialect/Triton/IR/Dialect.cpp.inc" -#include "triton/Dialect/Triton/IR/OpInterfaces.cpp.inc" - -using namespace mlir; -using namespace mlir::triton; - -//===----------------------------------------------------------------------===// -// TritonDialect Dialect Interfaces -//===----------------------------------------------------------------------===// - -bool TritonInlinerInterface::isLegalToInline(Operation *call, - Operation *callable, - bool wouldBeCloned) const { - auto funcOp = dyn_cast(callable); - if (!funcOp) - return true; - if (funcOp->hasAttr("noinline")) - return !funcOp->getAttrOfType("noinline").getValue(); - return true; -} - -/// Handle the given inlined terminator by replacing it with a new operation -/// as necessary. -void TritonInlinerInterface::handleTerminator(Operation *op, - Block *newDest) const { - // Only return needs to be handled here. - auto returnOp = dyn_cast(op); - if (!returnOp) - return; - - // Replace the return with a branch to the dest. - OpBuilder builder(op); - mlir::cf::BranchOp::create(builder, op->getLoc(), newDest, - returnOp.getOperands()); - op->erase(); -} - -/// Handle the given inlined terminator by replacing it with a new operation -/// as necessary. -void TritonInlinerInterface::handleTerminator(Operation *op, - ValueRange valuesToRepl) const { - // Only return needs to be handled here. - auto returnOp = cast(op); - - // Replace the values directly with the return operands. - assert(returnOp.getNumOperands() == valuesToRepl.size()); - for (const auto &it : llvm::enumerate(returnOp.getOperands())) - valuesToRepl[it.index()].replaceAllUsesWith(it.value()); -} - -void TritonDialect::initialize() { - registerTypes(); - - addOperations< -#define GET_OP_LIST -#include "triton/Dialect/Triton/IR/Ops.cpp.inc" - >(); - - // We can also add interface here. - addInterfaces(); -} - -Operation *TritonDialect::materializeConstant(OpBuilder &builder, - Attribute value, Type type, - Location loc) { - return arith::ConstantOp::materialize(builder, value, type, loc); -} diff --git a/third_party/mthreads/lib/Dialect/Triton/IR/DiscardableAttributes.cpp b/third_party/mthreads/lib/Dialect/Triton/IR/DiscardableAttributes.cpp deleted file mode 100644 index 8f4d80ea8a..0000000000 --- a/third_party/mthreads/lib/Dialect/Triton/IR/DiscardableAttributes.cpp +++ /dev/null @@ -1,17 +0,0 @@ -#include "mlir/Support/LLVM.h" -#include "triton/Dialect/Triton/IR/Dialect.h" - -namespace mlir::triton { - -SmallVector -filterDiscardableAttrs(Operation *op, ArrayRef allowList) { - SmallVector propagatedAttrs; - for (auto attrName : allowList) { - Attribute attr = op->getDiscardableAttr(attrName); - if (attr) - propagatedAttrs.emplace_back(attrName, attr); - } - return propagatedAttrs; -} - -} // namespace mlir::triton diff --git a/third_party/mthreads/lib/Dialect/Triton/IR/OpInterfaces.cpp b/third_party/mthreads/lib/Dialect/Triton/IR/OpInterfaces.cpp deleted file mode 100644 index 7bebffe61b..0000000000 --- a/third_party/mthreads/lib/Dialect/Triton/IR/OpInterfaces.cpp +++ /dev/null @@ -1,77 +0,0 @@ -#include "mlir/IR/BuiltinTypes.h" -#include "mlir/IR/Diagnostics.h" -#include "mlir/Support/LogicalResult.h" - -#include "triton/Dialect/Triton/IR/OpInterfaces.h" -#include "triton/Dialect/Triton/IR/Types.h" - -namespace mlir { -namespace triton { -namespace impl { - -LogicalResult verifyTransposeOpInterface(Operation *op) { - TransposeOpInterface transposeOp = cast(op); - auto rank = cast(transposeOp.getSrc().getType()).getRank(); - auto order = transposeOp.getOrder(); - if (static_cast(rank) != order.size()) { - return op->emitError( - "order must have the same size as the rank of the operand and result"); - } - - SmallVector sortedOrder(order); - llvm::sort(sortedOrder); - for (int32_t i = 0; i < sortedOrder.size(); i++) { - if (sortedOrder[i] != i) { - return op->emitError("order must be a permutation of [0, ..., rank - 1]"); - } - } - - return success(); -} - -// A DotOpInterface operation should have at least three operands. -// The first two operands should share a common dimension, and the result -// should have the dimensions of the two operands that are not shared. -// A DotOpInterface operation can be either 2d or 3d. -// In the 3d case, the first dimension of operands is the batch dimension. -LogicalResult verifyDotOpInterface(Operation *op) { - DotOpInterface dotOp = cast(op); - - if (dotOp->getNumOperands() < 3) - return dotOp->emitOpError("expected at least 3 operands"); - auto aTy = cast(dotOp->getOperand(0).getType()); - auto bTy = cast(dotOp->getOperand(1).getType()); - auto cTy = cast(dotOp->getOperand(2).getType()); - auto aShape = aTy.getShape(); - auto bShape = bTy.getShape(); - auto cShape = cTy.getShape(); - // Check if all 3d or all 2d - if (aShape.size() != 2 && aShape.size() != 3) - return dotOp->emitOpError("expected operands to be 2d or 3d"); - if (aShape.size() != bShape.size() || aShape.size() != cShape.size()) - return dotOp->emitOpError("expected all operands to have the same rank"); - - // Check for valid A, B input shapes for dot - if (!dotOp.verifyDims()) - return dotOp->emitOpError( - "expected the last dimension of the first operand " - "to be equal to the second-to-last dimension of " - "the second operand"); - - // Check the batch dimension - if (aShape.size() == 3 && (aShape[0] != cShape[0] || bShape[0] != cShape[0])) - return dotOp->emitOpError("expected the first dimension of the first " - "operand to be equal to the first dimension of " - "the result"); - // Check the output shape - if (!dotOp.verifyOutputDims()) - return dotOp->emitOpError( - "expected the output shape to be the concatenation of the last " - "dimension of the first operand and the last dimension of the " - "second "); - return success(); -} - -} // namespace impl -} // namespace triton -} // namespace mlir diff --git a/third_party/mthreads/lib/Dialect/Triton/Transforms/ArithTypeConversion.cpp b/third_party/mthreads/lib/Dialect/Triton/Transforms/ArithTypeConversion.cpp deleted file mode 100644 index 3928119409..0000000000 --- a/third_party/mthreads/lib/Dialect/Triton/Transforms/ArithTypeConversion.cpp +++ /dev/null @@ -1,51 +0,0 @@ -#include "triton/Dialect/Triton/Transforms/ArithTypeConversion.h" - -#include "mlir/Dialect/Arith/IR/Arith.h" -#include "mlir/Dialect/SCF/IR/SCF.h" -#include "mlir/IR/PatternMatch.h" -#include "mlir/Support/LLVM.h" -#include "mlir/Transforms/DialectConversion.h" - -namespace { - -struct RewriteArithSelectOp : mlir::OpConversionPattern { - using mlir::OpConversionPattern::OpConversionPattern; - - mlir::LogicalResult - matchAndRewrite(mlir::arith::SelectOp op, OneToNOpAdaptor adaptor, - mlir::ConversionPatternRewriter &rewriter) const override { - // Note we're replacing the select op with an if op because we are - // converting one value into many values. - auto newIf = mlir::scf::IfOp::create( - rewriter, op.getLoc(), mlir::TypeRange(adaptor.getTrueValue()), - op.getCondition(), true); - // We set the attributes from the op in case the op has any additional - // attributes - newIf->setAttrs(op->getAttrs()); - - { - mlir::ConversionPatternRewriter::InsertionGuard guard(rewriter); - rewriter.setInsertionPointToStart(newIf.thenBlock()); - mlir::scf::YieldOp::create(rewriter, op->getLoc(), - adaptor.getTrueValue()); - rewriter.setInsertionPointToStart(newIf.elseBlock()); - mlir::scf::YieldOp::create(rewriter, op->getLoc(), - adaptor.getFalseValue()); - } - - // Replace the old operation results - rewriter.replaceOpWithMultiple(op, {newIf->getResults()}); - - return mlir::success(); - } -}; - -} // namespace -namespace mlir::triton { - -void populateArithTypeConversions(const TypeConverter &converter, - RewritePatternSet &patterns) { - patterns.add(converter, patterns.getContext()); -} - -} // namespace mlir::triton diff --git a/third_party/mthreads/lib/Dialect/Triton/Transforms/CMakeLists.txt b/third_party/mthreads/lib/Dialect/Triton/Transforms/CMakeLists.txt deleted file mode 100644 index 8be846f589..0000000000 --- a/third_party/mthreads/lib/Dialect/Triton/Transforms/CMakeLists.txt +++ /dev/null @@ -1,27 +0,0 @@ -set(LLVM_TARGET_DEFINITIONS Combine.td) -mlir_tablegen(TritonCombine.inc -gen-rewriters) -add_public_tablegen_target(TritonCombineIncGen) - -add_triton_library(TritonTransforms - Combine.cpp - LoopAwareCSE.cpp - LoopInvariantCodeMotion.cpp - LoopPeeling.cpp - LoopUnroll.cpp - ReorderBroadcast.cpp - RewriteTensorPointer.cpp - RewriteTensorDescriptorToPointer.cpp - ArithTypeConversion.cpp - FunctionTypeConversion.cpp - - DEPENDS - TritonTransformsIncGen - TritonCombineIncGen - - LINK_LIBS PUBLIC - MLIRPass - MLIRTransformUtils - MLIRTransforms - MLIRSCFToControlFlow - TritonIR -) diff --git a/third_party/mthreads/lib/Dialect/Triton/Transforms/FunctionTypeConversion.cpp b/third_party/mthreads/lib/Dialect/Triton/Transforms/FunctionTypeConversion.cpp deleted file mode 100644 index f3a454abea..0000000000 --- a/third_party/mthreads/lib/Dialect/Triton/Transforms/FunctionTypeConversion.cpp +++ /dev/null @@ -1,163 +0,0 @@ -#include "triton/Dialect/Triton/Transforms/FunctionTypeConversion.h" - -#include "mlir/IR/Value.h" -#include "mlir/Support/LLVM.h" -#include "mlir/Transforms/DialectConversion.h" -#include "triton/Dialect/Triton/IR/Dialect.h" -#include "llvm/ADT/STLExtras.h" -#include "llvm/ADT/SmallVector.h" - -#include - -namespace mlir::triton { - -namespace { - -SmallVector flattenValues(ArrayRef values) { - SmallVector ret; - for (const auto &vs : values) { - llvm::append_range(ret, vs); - } - return ret; -} - -struct CallOpConversion : public OpConversionPattern { - using OpConversionPattern::OpConversionPattern; - - LogicalResult - matchAndRewrite(CallOp callOp, OneToNOpAdaptor adaptor, - ConversionPatternRewriter &rewriter) const override { - llvm::SmallVector resultReplacementGrouping; - llvm::SmallVector convertedResults; - - for (auto type : callOp->getResultTypes()) { - const auto oldNumFlattenedResults = convertedResults.size(); - if (failed(getTypeConverter()->convertTypes(type, convertedResults))) { - return failure(); - } - resultReplacementGrouping.push_back(convertedResults.size() - - oldNumFlattenedResults); - } - - auto newCallOp = - CallOp::create(rewriter, callOp->getLoc(), callOp.getCallee(), - convertedResults, flattenValues(adaptor.getOperands())); - // Preserve any additional attributes that may have been set on the op - newCallOp->setAttrs(callOp->getAttrs()); - - SmallVector replacements; - std::size_t offset = 0; - for (auto groupSize : resultReplacementGrouping) { - replacements.push_back(newCallOp->getResults().slice(offset, groupSize)); - offset += groupSize; - } - - rewriter.replaceOpWithMultiple(callOp, replacements); - return success(); - } -}; - -struct ReturnOpConversion : public OpConversionPattern { - using OpConversionPattern::OpConversionPattern; - - LogicalResult - matchAndRewrite(ReturnOp returnOp, OneToNOpAdaptor adaptor, - ConversionPatternRewriter &rewriter) const override { - auto newReturnOp = ReturnOp::create(rewriter, returnOp->getLoc(), - flattenValues(adaptor.getOperands())); - // Preserve any additional attributes that may have been set on the op - newReturnOp->setAttrs(returnOp->getAttrs()); - - rewriter.replaceOp(returnOp, newReturnOp); - return success(); - } -}; - -//===----------------------------------------------------------------------===// -// FunctionOpInterfaceSignatureConversion -//===----------------------------------------------------------------------===// -// NOTE: Forked from mlir to support remapping argument attributes correctly in -// a one-to-many type conversion. - -SmallVector -convertFuncOpAttrs(FunctionOpInterface funcOp, - TypeConverter::SignatureConversion &sigConv, - FunctionType newType) { - if (newType.getNumInputs() == funcOp.getNumArguments()) { - return {}; - } - ArrayAttr allArgAttrs = funcOp.getAllArgAttrs(); - if (!allArgAttrs) - return {}; - - SmallVector newAttrs(newType.getNumInputs()); - for (auto i : llvm::seq(allArgAttrs.size())) { - auto mapping = sigConv.getInputMapping(i); - assert(mapping.has_value()); - auto outIdx = mapping->inputNo; - newAttrs[outIdx] = allArgAttrs[i]; - } - return newAttrs; -} - -LogicalResult convertFuncOpTypes(FunctionOpInterface funcOp, - const TypeConverter &typeConverter, - ConversionPatternRewriter &rewriter) { - FunctionType type = dyn_cast(funcOp.getFunctionType()); - if (!type) - return failure(); - - // Convert the original function types. - TypeConverter::SignatureConversion result(type.getNumInputs()); - SmallVector newResults; - if (failed(typeConverter.convertSignatureArgs(type.getInputs(), result)) || - failed(typeConverter.convertTypes(type.getResults(), newResults)) || - failed(rewriter.convertRegionTypes(&funcOp.getFunctionBody(), - typeConverter, &result))) - return failure(); - - // Update the function signature in-place. - auto newType = FunctionType::get(rewriter.getContext(), - result.getConvertedTypes(), newResults); - - auto newArgAttrs = convertFuncOpAttrs(funcOp, result, newType); - - rewriter.modifyOpInPlace(funcOp, [&] { - funcOp.setType(newType); - if (!newArgAttrs.empty()) { - funcOp.setAllArgAttrs(newArgAttrs); - } - }); - - return success(); -} - -/// Create a default conversion pattern that rewrites the type signature of a -/// FunctionOpInterface op. This only supports ops which use FunctionType to -/// represent their type. -struct FunctionOpInterfaceSignatureConversion : public ConversionPattern { - FunctionOpInterfaceSignatureConversion(StringRef functionLikeOpName, - MLIRContext *ctx, - const TypeConverter &converter, - PatternBenefit benefit = 1) - : ConversionPattern(converter, functionLikeOpName, benefit, ctx) {} - - LogicalResult - matchAndRewrite(Operation *op, ArrayRef /*operands*/, - ConversionPatternRewriter &rewriter) const override { - FunctionOpInterface funcOp = cast(op); - return convertFuncOpTypes(funcOp, *typeConverter, rewriter); - } -}; - -} // namespace - -void populateFunctionTypeConversions(const TypeConverter &converter, - RewritePatternSet &patterns) { - auto context = patterns.getContext(); - patterns.add( - triton::FuncOp::getOperationName(), context, converter); - patterns.add(converter, context); -} - -} // namespace mlir::triton diff --git a/third_party/mthreads/lib/Dialect/Triton/Transforms/LoopAwareCSE.cpp b/third_party/mthreads/lib/Dialect/Triton/Transforms/LoopAwareCSE.cpp deleted file mode 100644 index ad9ca7f396..0000000000 --- a/third_party/mthreads/lib/Dialect/Triton/Transforms/LoopAwareCSE.cpp +++ /dev/null @@ -1,178 +0,0 @@ -#include "mlir/Dialect/SCF/IR/SCF.h" -#include "mlir/IR/Dominance.h" -#include "mlir/Pass/Pass.h" -#include "mlir/Transforms/CSE.h" -#include "mlir/Transforms/GreedyPatternRewriteDriver.h" -#include "llvm/ADT/EquivalenceClasses.h" - -using namespace mlir; - -namespace mlir::triton { -#define GEN_PASS_DEF_TRITONLOOPAWARECSE -#include "triton/Dialect/Triton/Transforms/Passes.h.inc" -} // namespace mlir::triton - -namespace { -class ValueEquivalence { -public: - std::optional getKnownEquivalence(Value a, Value b) { - if (auto it = equalValues.find(normalizeKey(a, b)); it != equalValues.end()) - return it->second; - return std::nullopt; - } - void setKnownEquivalence(Value a, Value b, bool eq) { - equalValues.insert_or_assign(normalizeKey(a, b), eq); - } - -private: - // Commutatively query the equivalence of two values by sorting the key by - // pointer value. - std::pair normalizeKey(Value a, Value b) { - if ((uintptr_t)a.getAsOpaquePointer() < (uintptr_t)b.getAsOpaquePointer()) - return {a, b}; - return {b, a}; - } - - DenseMap, bool> equalValues; -}; - -struct LoopCSEDriver { - LoopCSEDriver(scf::ForOp loop) : loop(loop) {} - - bool areIterArgsEqual(int i, int j); - bool areEqualInLoop(Value a, Value b); - - scf::ForOp loop; - SmallVector> argStack; -}; -} // namespace - -bool LoopCSEDriver::areIterArgsEqual(int i, int j) { - if (i == j) - return true; - if (loop.getInitArgs()[i] != loop.getInitArgs()[j]) - return false; - if (llvm::is_contained(argStack, std::make_pair(i, j))) - return true; - - // First, assume the arguments are equal. This is how recursion is broken. - argStack.push_back({i, j}); - bool result = - areEqualInLoop(loop.getYieldedValues()[i], loop.getYieldedValues()[j]); - argStack.pop_back(); - return result; -} - -bool LoopCSEDriver::areEqualInLoop(Value a, Value b) { - // Check trivial case. - if (a == b) - return true; - if (a.getType() != b.getType()) - return false; - - Block *aBlock = a.getParentBlock(); - Block *bBlock = b.getParentBlock(); - // Values from outside the loop must have been equal. - if (aBlock != loop.getBody() || bBlock != loop.getBody()) { - return false; - } - // Both must be block arguments or not. - if (isa(a) != isa(b)) - return false; - // Both must be the inductor var or not. - if (a == loop.getInductionVar() || b == loop.getInductionVar()) - return false; - - if (auto aArg = dyn_cast(a)) { - auto bArg = cast(b); - bool result = - areIterArgsEqual(aArg.getArgNumber() - 1, bArg.getArgNumber() - 1); - return result; - } - - Operation *aDef = a.getDefiningOp(); - Operation *bDef = b.getDefiningOp(); - if (cast(a).getResultNumber() != - cast(b).getResultNumber()) - return false; - // For it to be known that the operation results have the same value, they - // must be side effect free. - if (!isMemoryEffectFree(aDef) || !isMemoryEffectFree(bDef)) - return false; - // Don't bother with operations with regions. - if (aDef->getNumRegions() || bDef->getNumRegions()) - return false; - - bool result = OperationEquivalence::isEquivalentTo( - aDef, bDef, - [&](Value a, Value b) { return success(areEqualInLoop(a, b)); }, - /*markEquivalent=*/nullptr, OperationEquivalence::IgnoreLocations); - return result; -} - -static void loopCSE(scf::ForOp loop) { - int numIterArgs = loop.getNumRegionIterArgs(); - // Group equivalent iter args together. - llvm::EquivalenceClasses equivalentArgs; - LoopCSEDriver driver(loop); - for (int i = 0; i != numIterArgs; ++i) { - for (int j = i + 1; j != numIterArgs; ++j) { - if (driver.areIterArgsEqual(i, j)) - equivalentArgs.unionSets(i, j); - } - } - - // For each equivalence class, replace all other args in the class with one. - for (auto it = equivalentArgs.begin(), end = equivalentArgs.end(); it != end; - ++it) { - if (!(*it)->isLeader()) - continue; - SmallVector eqArgs; - for (auto mIt = equivalentArgs.member_begin(**it); - mIt != equivalentArgs.member_end(); ++mIt) - eqArgs.push_back(*mIt); - assert(eqArgs.size() > 1); - // Sort the indices so the pass is deterministic. - llvm::sort(eqArgs); - BlockArgument unique = loop.getRegionIterArg(eqArgs.front()); - Value uniqueResult = loop.getResult(eqArgs.front()); - for (int j : llvm::drop_begin(eqArgs)) { - BlockArgument other = loop.getRegionIterArg(j); - other.replaceAllUsesWith(unique); - // Short-circuit the value. The canonicalizer will clean this up. Leftover - // subcomputations can now be removed by normal CSE. - (*loop.getYieldedValuesMutable())[j].set(other); - loop.getResult(j).replaceAllUsesWith(uniqueResult); - } - } -} - -namespace { -struct LoopAwareCSE - : public triton::impl::TritonLoopAwareCSEBase { - using TritonLoopAwareCSEBase::TritonLoopAwareCSEBase; - - void runOnOperation() override { - // LoopAwareCSE doesn't recursively CSE ops outside of loops, so run CSE - // first to make sure values from outside loops that are equivalent are made - // pointer equal. - IRRewriter rewriter(&getContext()); - auto &domInfo = getAnalysis(); - eliminateCommonSubExpressions(rewriter, domInfo, getOperation()); - - // CSE region iter args within loop bodies. - getOperation().walk(loopCSE); - - // Now that equivalent iter args have been made pointer equal, run CSE again - // to clean up the loop body. - eliminateCommonSubExpressions(rewriter, domInfo, getOperation()); - - // Run the `scf.for` canonicalizer to clean up the loops (short-circuited - // values, unused results, etc.). - RewritePatternSet patterns(&getContext()); - scf::ForOp::getCanonicalizationPatterns(patterns, &getContext()); - if (failed(applyPatternsGreedily(getOperation(), std::move(patterns)))) - return signalPassFailure(); - } -}; -} // namespace diff --git a/third_party/mthreads/lib/Dialect/Triton/Transforms/LoopInvariantCodeMotion.cpp b/third_party/mthreads/lib/Dialect/Triton/Transforms/LoopInvariantCodeMotion.cpp deleted file mode 100644 index a1de3bf845..0000000000 --- a/third_party/mthreads/lib/Dialect/Triton/Transforms/LoopInvariantCodeMotion.cpp +++ /dev/null @@ -1,82 +0,0 @@ -#include "mlir/Transforms/LoopInvariantCodeMotionUtils.h" -#include "triton/Dialect/Triton/IR/Dialect.h" -#include "triton/Dialect/Triton/IR/Utility.h" -#include "triton/Dialect/Triton/Transforms/Passes.h" -#include "llvm/Support/Debug.h" - -namespace mlir::triton { - -#define GEN_PASS_DEF_TRITONLOOPINVARIANTCODEMOTION -#include "triton/Dialect/Triton/Transforms/Passes.h.inc" - -#define DEBUG_TYPE "triton-licm" -#define DBGS() (llvm::dbgs() << "[" DEBUG_TYPE "]: ") -#define LDBG(X) LLVM_DEBUG(DBGS() << X << "\n") - -class LoopInvariantCodeMotionPass - : public impl::TritonLoopInvariantCodeMotionBase< - LoopInvariantCodeMotionPass> { - - DenseMap isLoopMemoryEffectFreeOrOnlyRead; - - bool isMemoryEffectFreeOrOnlyRead(Operation *op) { - std::optional> effects = - getEffectsRecursively(op); - if (!effects) - return false; - return llvm::all_of(*effects, - [&](const MemoryEffects::EffectInstance &effect) { - return isa(effect.getEffect()); - }); - } - - void runOnOperation() override { - // Walk through all loops in a function in innermost-loop-first order. - // This way, we first LICM from the inner loop, and place the ops in the - // outer loop, which in turn can be further LICM'ed. - getOperation()->walk([&](LoopLikeOpInterface loopLike) { - moveLoopInvariantCode( - loopLike.getLoopRegions(), - // isDefinedOutsideOfRegion - [&](Value value, Region *region) { - return loopLike.isDefinedOutsideOfLoop(value); - }, - // shouldMoveOutOfRegion - [&](Operation *op, Region *region) { - if (!isa(op)) - return isSpeculatable(op) && isMemoryEffectFree(op); - if (!isLoopMemoryEffectFreeOrOnlyRead.contains(loopLike)) - isLoopMemoryEffectFreeOrOnlyRead[loopLike] = - isMemoryEffectFreeOrOnlyRead(loopLike); - return isMemoryEffectFreeOrOnlyRead(op) && - isLoopMemoryEffectFreeOrOnlyRead[loopLike]; - }, - // moveOutOfRegion - [&](Operation *op, Region *) { - // Create the new mask for load op. - if (auto loadOp = dyn_cast(op)) { - IRRewriter rewriter(loopLike); - Location loc = loopLike->getLoc(); - Value cond; - if (auto forOp = dyn_cast(loopLike.getOperation())) { - cond = arith::CmpIOp::create( - rewriter, loc, arith::CmpIPredicate::slt, - forOp.getLowerBound(), forOp.getUpperBound()); - } else if (auto whileOp = - dyn_cast(loopLike.getOperation())) { - // TODO: Support Load Op hoisting for while loop. - return; - } else { - return; - } - Value newMask = getPredMask(rewriter, loadOp.getPtr().getType(), - loadOp.getMask(), cond); - loadOp.getMaskMutable().assign(newMask); - } - loopLike.moveOutOfLoop(op); - }); - }); - } -}; - -} // namespace mlir::triton diff --git a/third_party/mthreads/lib/Dialect/Triton/Transforms/LoopPeeling.cpp b/third_party/mthreads/lib/Dialect/Triton/Transforms/LoopPeeling.cpp deleted file mode 100644 index ed887bfee0..0000000000 --- a/third_party/mthreads/lib/Dialect/Triton/Transforms/LoopPeeling.cpp +++ /dev/null @@ -1,67 +0,0 @@ -#include "triton/Dialect/Triton/Transforms/LoopPeeling.h" -#include "mlir/Dialect/SCF/IR/SCF.h" -#include "mlir/Pass/Pass.h" -#include "triton/Dialect/Triton/IR/Utility.h" - -using namespace mlir; - -namespace mlir { -namespace triton { - -void peelLoopEpilogue( - scf::ForOp forOp, - function_ref - processPeeledOp) { - SmallVector loopBodyOps; - IRRewriter rewriter(forOp); - Location loc = forOp.getLoc(); - Type type = forOp.getStep().getType(); - - // Fetch loop bounds and step - Value lowerBound = forOp.getLowerBound(); - Value upperBound = forOp.getUpperBound(); - Value step = forOp.getStep(); - Value newUpperBound = arith::SubIOp::create(rewriter, loc, upperBound, step); - - rewriter.setInsertionPointAfter(forOp); - Value lastIV = getLastInductionValue(rewriter, forOp); - - auto cond = arith::CmpIOp::create(rewriter, loc, arith::CmpIPredicate::slt, - lowerBound, upperBound); - - // Create an if op to execute the peeled iteration - IRMapping map; - map.map(forOp.getRegionIterArgs(), forOp.getResults()); - map.map(forOp.getInductionVar(), lastIV); - auto ifOp = scf::IfOp::create(rewriter, loc, forOp.getResultTypes(), cond); - forOp.getBodyRegion().cloneInto(&ifOp.getThenRegion(), map); - auto newElseBlock = rewriter.createBlock(&ifOp.getElseRegion()); - rewriter.setInsertionPointToStart(newElseBlock); - scf::YieldOp::create(rewriter, loc, forOp.getResults()); - - forOp->replaceUsesWithIf(ifOp, [&](OpOperand &operand) { - return !ifOp->isAncestor(operand.getOwner()); - }); - - forOp.getUpperBoundMutable().assign(newUpperBound); - - if (processPeeledOp) { - for (auto &op : - llvm::make_early_inc_range(forOp.getBody()->without_terminator())) { - Operation *newOp = processPeeledOp(rewriter, &op, /*isEpilogue=*/false); - if (newOp && newOp != &op) { - op.replaceAllUsesWith(newOp); - } - } - for (auto &op : llvm::make_early_inc_range( - ifOp.getThenRegion().front().without_terminator())) { - Operation *newOp = processPeeledOp(rewriter, &op, /*isEpilogue=*/true); - if (newOp && newOp != &op) { - op.replaceAllUsesWith(newOp); - } - } - } -} - -} // namespace triton -} // namespace mlir diff --git a/third_party/mthreads/lib/Dialect/Triton/Transforms/ReorderBroadcast.cpp b/third_party/mthreads/lib/Dialect/Triton/Transforms/ReorderBroadcast.cpp deleted file mode 100644 index bdb8e527f9..0000000000 --- a/third_party/mthreads/lib/Dialect/Triton/Transforms/ReorderBroadcast.cpp +++ /dev/null @@ -1,230 +0,0 @@ -#include - -#include "mlir/IR/BuiltinAttributes.h" -#include "mlir/IR/Matchers.h" -#include "mlir/IR/PatternMatch.h" -#include "mlir/Pass/Pass.h" -#include "mlir/Support/LLVM.h" -#include "mlir/Support/LogicalResult.h" -#include "mlir/Transforms/GreedyPatternRewriteDriver.h" -#include "triton/Dialect/Triton/IR/Dialect.h" -#include "triton/Dialect/Triton/Transforms/Passes.h" - -namespace mlir::triton { - -#define GEN_PASS_DEF_TRITONREORDERBROADCAST -#include "triton/Dialect/Triton/Transforms/Passes.h.inc" - -namespace { - -Operation *cloneWithNewArgsAndResultTypes(PatternRewriter &rewriter, - Operation *op, ValueRange newOperands, - TypeRange newTypes) { - OperationState newElementwiseState(op->getLoc(), op->getName()); - newElementwiseState.addOperands(newOperands); - newElementwiseState.addTypes(newTypes); - newElementwiseState.addAttributes(op->getAttrs()); - return rewriter.create(newElementwiseState); -} - -bool isSplat(Operation *op) { - if (auto splatOp = llvm::dyn_cast(op)) { - return true; - } - DenseElementsAttr constAttr; - return (matchPattern(op, m_Constant(&constAttr)) && constAttr.isSplat()); -} - -// elementwise(splat(a), splat(b), ...) => splat(elementwise(a, b, ...)) -struct MoveSplatAfterElementwisePattern - : public OpTraitRewritePattern { - - MoveSplatAfterElementwisePattern(MLIRContext *context) - : OpTraitRewritePattern(context) {} - - LogicalResult matchAndRewrite(Operation *op, - PatternRewriter &rewriter) const override { - if (!isMemoryEffectFree(op)) { - return failure(); - } - - for (auto operand : op->getOperands()) { - auto definingOp = operand.getDefiningOp(); - if (!definingOp) - return failure(); - - if (!isSplat(definingOp)) { - return failure(); - } - } - - if (op->getNumOperands() <= 0) - return failure(); - - auto loc = op->getLoc(); - auto operands = op->getOperands(); - - llvm::SmallVector scalarOperands(operands.size()); - for (unsigned iOp = 0; iOp < operands.size(); ++iOp) { - auto definingOp = operands[iOp].getDefiningOp(); - - DenseElementsAttr constAttr; - if (auto splatOp = llvm::dyn_cast(definingOp)) { - scalarOperands[iOp] = splatOp.getSrc(); - } else if (matchPattern(definingOp, m_Constant(&constAttr)) && - constAttr.isSplat()) { - auto value = constAttr.getSplatValue(); - scalarOperands[iOp] = arith::ConstantOp::materialize( - rewriter, value, constAttr.getElementType(), loc); - } else { - llvm_unreachable("Expected a splat"); - } - } - - auto resultTypes = op->getResultTypes(); - llvm::SmallVector scalarResultTys; - for (auto resultTy : resultTypes) { - auto elemTy = dyn_cast(resultTy).getElementType(); - scalarResultTys.push_back(elemTy); - } - - auto newOp = cloneWithNewArgsAndResultTypes(rewriter, op, scalarOperands, - scalarResultTys); - - for (unsigned iRes = 0; iRes < resultTypes.size(); ++iRes) { - auto newResult = SplatOp::create(rewriter, loc, resultTypes[iRes], - newOp->getResult(iRes)); - rewriter.replaceAllUsesWith(op->getResult(iRes), newResult); - } - return success(); - } -}; - -// elementwise(broadcast(a)) => broadcast(elementwise(a)) -// This also generalizes to multiple arguments when the rest are splat-like -// Not handled: multiple broadcasted arguments -struct MoveBroadcastAfterElementwisePattern - : public OpTraitRewritePattern { - - MoveBroadcastAfterElementwisePattern(MLIRContext *context) - : OpTraitRewritePattern(context) {} - - LogicalResult matchAndRewrite(Operation *op, - PatternRewriter &rewriter) const override { - if (!isMemoryEffectFree(op)) { - return failure(); - } - - auto operands = op->getOperands(); - bool seenBroadcast = false; - ArrayRef srcShape; - for (auto operand : operands) { - auto definingOp = operand.getDefiningOp(); - if (!definingOp) { - return failure(); - } - auto getSrcShape = [](BroadcastOp b) { - return b.getSrc().getType().getShape(); - }; - if (auto broadcastOp = llvm::dyn_cast(definingOp)) { - if (!seenBroadcast) { - seenBroadcast = true; - srcShape = getSrcShape(broadcastOp); - } else if (srcShape != getSrcShape(broadcastOp)) { - // If the broadcast have different types we cannot re-order. - return failure(); - } - } else if (!isSplat(definingOp)) { - // Not splat or broadcast - return failure(); - } - } - if (!seenBroadcast) - return failure(); - - auto loc = op->getLoc(); - - // Find broadcast op - BroadcastOp broadcastOp; - for (auto operand : operands) { - broadcastOp = operand.getDefiningOp(); - if (broadcastOp) { - break; - } - } - - auto srcTy = broadcastOp.getSrc().getType(); - auto bcSrcShape = srcTy.getShape(); - - // Reshape operands to match srcShape - llvm::SmallVector newOperands; - for (auto operand : operands) { - auto definingOp = operand.getDefiningOp(); - if (auto broadcastSrcOp = llvm::dyn_cast(definingOp)) { - newOperands.push_back(broadcastSrcOp.getSrc()); - continue; - } - auto elemTy = - dyn_cast(operand.getType()).getElementType(); - auto newTy = srcTy.clone(bcSrcShape, elemTy); - if (auto splatOp = llvm::dyn_cast(definingOp)) { - auto newSplat = SplatOp::create(rewriter, loc, newTy, splatOp.getSrc()); - newOperands.push_back(newSplat); - continue; - } - DenseElementsAttr constAttr; - if (matchPattern(definingOp, m_Constant(&constAttr)) && - constAttr.isSplat()) { - auto scalarValue = constAttr.getSplatValue(); - auto splatValue = SplatElementsAttr::get(newTy, scalarValue); - auto newConstant = - arith::ConstantOp::create(rewriter, loc, newTy, splatValue); - newOperands.push_back(newConstant); - continue; - } - llvm_unreachable("Expected broadcast or splat"); - } - - // Reshape results to match srcShape - llvm::SmallVector newResultTypes; - auto resultTypes = op->getResultTypes(); - for (auto resultTy : resultTypes) { - auto elemTy = dyn_cast(resultTy).getElementType(); - newResultTypes.push_back(srcTy.clone(bcSrcShape, elemTy)); - } - - // Create new op and broadcast results - auto newOp = cloneWithNewArgsAndResultTypes(rewriter, op, newOperands, - newResultTypes); - for (unsigned iRes = 0; iRes < newResultTypes.size(); ++iRes) { - auto newResult = BroadcastOp::create(rewriter, loc, resultTypes[iRes], - newOp->getResult(iRes)); - rewriter.replaceAllUsesWith(op->getResult(iRes), newResult); - } - return success(); - } -}; - -} // namespace - -class ReorderBroadcastPass - : public impl::TritonReorderBroadcastBase { -public: - void runOnOperation() override { - MLIRContext *context = &getContext(); - RewritePatternSet patterns(context); - ModuleOp m = getOperation(); - - BroadcastOp::getCanonicalizationPatterns(patterns, context); - ExpandDimsOp::getCanonicalizationPatterns(patterns, context); - // elementwise(broadcast(a)) => broadcast(elementwise(a)) - patterns.add(context); - // elementwise(splat(a), splat(b), ...) => splat(elementwise(a, b, ...)) - patterns.add(context); - - if (applyPatternsGreedily(m, std::move(patterns)).failed()) - signalPassFailure(); - } -}; - -} // namespace mlir::triton diff --git a/third_party/mthreads/lib/Dialect/TritonGPU/CMakeLists.txt b/third_party/mthreads/lib/Dialect/TritonGPU/CMakeLists.txt deleted file mode 100644 index 9f57627c32..0000000000 --- a/third_party/mthreads/lib/Dialect/TritonGPU/CMakeLists.txt +++ /dev/null @@ -1,2 +0,0 @@ -add_subdirectory(IR) -add_subdirectory(Transforms) diff --git a/third_party/mthreads/lib/Dialect/TritonGPU/IR/CMakeLists.txt b/third_party/mthreads/lib/Dialect/TritonGPU/IR/CMakeLists.txt deleted file mode 100644 index 782b66b686..0000000000 --- a/third_party/mthreads/lib/Dialect/TritonGPU/IR/CMakeLists.txt +++ /dev/null @@ -1,18 +0,0 @@ -add_triton_library(TritonGPUIR - Dialect.cpp - LinearLayoutConversions.cpp - Ops.cpp - Types.cpp - - DEPENDS - TritonGPUCGAAttrIncGen - TritonGPUTableGen - TritonGPUAttrDefsIncGen - TritonGPUTypeInterfacesIncGen - TritonGPUOpInterfacesIncGen - - LINK_LIBS PUBLIC - MLIRGPUDialect - TritonIR - TritonTools -) diff --git a/third_party/mthreads/lib/Dialect/TritonGPU/Transforms/CMakeLists.txt b/third_party/mthreads/lib/Dialect/TritonGPU/Transforms/CMakeLists.txt deleted file mode 100644 index 962ef6a1fe..0000000000 --- a/third_party/mthreads/lib/Dialect/TritonGPU/Transforms/CMakeLists.txt +++ /dev/null @@ -1,56 +0,0 @@ -add_triton_library(TritonGPUTransforms - AccelerateMatmul.cpp - Coalesce.cpp - F32DotTC.cpp - FuseNestedLoops.cpp - CombineTensorSelectAndIf.cpp - DecomposeScaledBlocked.cpp - HoistTMEMAlloc.cpp - ReduceDataDuplication.cpp - OptimizeAccumulatorInit.cpp - OptimizeDotOperands.cpp - OptimizeThreadLocality.cpp - Pipeliner/AssignLatencies.cpp - Pipeliner/LowerLoops.cpp - Pipeliner/MMAv5PipelineUtility.cpp - Pipeliner/ScheduleLoops.cpp - Pipeliner/WGMMAPipeline.cpp - Pipeliner/PipelineExpander.cpp - Pipeliner/TestPipelineLowerLoop.cpp - Pipeliner/SoftwarePipeliner.cpp - Pipeliner/TMAStoresPipeline.cpp - Pipeliner/MMAv5PipelineUtility.cpp - Pipeliner/PipeliningUtility.cpp - Pipeliner/Schedule.cpp - Prefetch.cpp - RemoveLayoutConversions.cpp - ReorderInstructions.cpp - CoalesceAsyncCopy.cpp - Utility.cpp - CoalesceUtils.cpp - LayoutPropagationUtility.cpp - WarpSpecialization/AutomaticWarpSpecialization.cpp - WarpSpecialization/Partition.cpp - WarpSpecialization/OptimizePartitionWarps.cpp - WarpSpecialization/PartitionBuilder.cpp - WarpSpecialization/PartitionLoops.cpp - WarpSpecialization/PartitionScheduling.cpp - WarpSpecialization/PartitionSchedulingUtility.cpp - - DEPENDS - TritonGPUTransformsIncGen - - LINK_LIBS PUBLIC - MLIRTransforms - MLIRTransformUtils - TritonAnalysis - TritonIR - TritonTransforms - TritonGPUIR - TritonNvidiaGPUIR - NVWSIR - NVWSTransforms - TritonToTritonGPU - TritonInstrumentIR - MLIRTransformUtils -) diff --git a/third_party/mthreads/lib/Dialect/TritonGPU/Transforms/CombineTensorSelectAndIf.cpp b/third_party/mthreads/lib/Dialect/TritonGPU/Transforms/CombineTensorSelectAndIf.cpp deleted file mode 100644 index 608e65a153..0000000000 --- a/third_party/mthreads/lib/Dialect/TritonGPU/Transforms/CombineTensorSelectAndIf.cpp +++ /dev/null @@ -1,176 +0,0 @@ -#include "mlir/Analysis/TopologicalSortUtils.h" -#include "mlir/IR/Dominance.h" -#include "mlir/Support/LLVM.h" -#include "mlir/Transforms/Passes.h" -#include "triton/Analysis/Utility.h" -#include "triton/Dialect/TritonGPU/IR/Dialect.h" -#include "triton/Dialect/TritonGPU/Transforms/Passes.h" -#include "triton/Dialect/TritonGPU/Transforms/Utility.h" - -#include - -namespace mlir { -namespace triton { -namespace gpu { - -#define GEN_PASS_DEF_TRITONGPUCOMBINETENSORSELECTANDIF -#include "triton/Dialect/TritonGPU/Transforms/Passes.h.inc" - -/// The user of select maybe inside either the ThenRegion or ElseRegion of -/// the scf.if. So, canonicalize user of select in scf.if first. -static void canonicalizeSelectUsersInSCFIf(ModuleOp input) { - llvm::MapVector, SmallVector> - usersNeedreplaced; - input.walk([&](arith::SelectOp selectOp) { - auto *parentBlock = selectOp->getBlock(); - Value condition = selectOp.getOperand(0); - Value trueVal = selectOp.getOperand(1); - Value falseVal = selectOp.getOperand(2); - Value resVal = selectOp.getResult(); - for (auto *condUser : condition.getUsers()) { - if (!llvm::isa(condUser)) - continue; - scf::IfOp ifOp = llvm::cast(condUser); - for (auto *resUser : resVal.getUsers()) { - if (ifOp->isProperAncestor(resUser)) { - if (ifOp.getThenRegion().findAncestorOpInRegion(*resUser) != - nullptr) { - // The user is inside the ThenRegion of the scf.if. - usersNeedreplaced[std::make_pair(resVal, trueVal)].push_back( - resUser); - } else { - // The user is inside the ElseRegion of the scf.if. - usersNeedreplaced[std::make_pair(resVal, falseVal)].push_back( - resUser); - } - } - } - } - }); - - // Replace the operand of user. - for (auto [replacedSrcAndDst, users] : - llvm::make_early_inc_range(usersNeedreplaced)) { - Value srcVal = replacedSrcAndDst.first; - Value dstVal = replacedSrcAndDst.second; - for (Operation *user : llvm::make_early_inc_range(users)) { - srcVal.replaceUsesWithIf( - dstVal, [&](OpOperand &use) { return use.getOwner() == user; }); - } - } -} - -/// Return true if the select could be merged into the If without breaking SSA -/// rules. -static bool canMergeIntoIf(arith::SelectOp selectOp, scf::IfOp ifOp, - DominanceInfo &dom) { - // If needs to be dominated by the select. - if (!dom.dominates(selectOp.getOperation(), ifOp.getOperation())) { - return false; - } - // If needs to dominate all the select's users. - for (auto user : selectOp.getResult().getUsers()) { - if (!dom.dominates(ifOp, user)) { - return false; - } - } - return true; -} - -class CombineTensorSelectAndIfPass - : public impl::TritonGPUCombineTensorSelectAndIfBase< - CombineTensorSelectAndIfPass> { -public: - void runOnOperation() override { - MLIRContext *context = &getContext(); - ModuleOp m = getOperation(); - canonicalizeSelectUsersInSCFIf(m); - - // Go over the arith.select ops, look if there is an if - // with the same condition. - DominanceInfo dom(m); - llvm::MapVector> selectToIf; - m.walk([&](arith::SelectOp selectOp) { - // Apply only to selects with a tensor result. Scalars are cheap enough to - // predicate. - if (!isa(selectOp.getResult().getType())) - return; - // Look if there is an if in the same block, with the same condition. - auto *parentBlock = selectOp->getBlock(); - Value condition = selectOp.getOperand(0); - SetVector conditionUsers(condition.getUsers().begin(), - condition.getUsers().end()); - // sort the users in topological order. - conditionUsers = mlir::topologicalSort(conditionUsers); - // Get condition's users - for (Operation *user : conditionUsers) { - auto ifOp = dyn_cast(user); - if (!ifOp || ifOp->getBlock() != parentBlock) - continue; - if (canMergeIntoIf(selectOp, ifOp, dom)) { - selectToIf[ifOp].push_back(selectOp); - break; - } - } - }); - - for (auto [ifOp, selectOps] : selectToIf) { - // Add new return value to the if (and create else block if necessary), - // then yield the select value in the then block and the else block. - OpBuilder builder(ifOp); - auto loc = ifOp.getLoc(); - // Create an scf::IfOp with extra return value. - SmallVector newResultTypes = {ifOp.getResultTypes().begin(), - ifOp.getResultTypes().end()}; - for (arith::SelectOp selectOp : selectOps) { - newResultTypes.push_back(selectOp.getResult().getType()); - } - auto newIfOp = scf::IfOp::create(builder, loc, newResultTypes, - ifOp.getCondition(), /*hasElse*/ true); - // Move the existing blocks to the new if. - newIfOp.getThenRegion().takeBody(ifOp.getThenRegion()); - - if (ifOp.elseBlock()) { - newIfOp.getElseRegion().takeBody(ifOp.getElseRegion()); - } else { - // Create an empty yield - auto builder = newIfOp.getElseBodyBuilder(); - auto yieldOp = scf::YieldOp::create(builder, loc); - } - - SmallVector ifYieldOperands = newIfOp.thenYield().getOperands(); - SmallVector elseYieldOperands = newIfOp.elseYield().getOperands(); - for (arith::SelectOp selectOp : selectOps) { - Value thenValue = selectOp.getTrueValue(); - Value elseValue = selectOp.getFalseValue(); - ifYieldOperands.push_back(thenValue); - elseYieldOperands.push_back(elseValue); - } - // Update yields - auto updateYield = [&](scf::YieldOp yield, SmallVector &operands) { - builder.setInsertionPoint(yield); - scf::YieldOp::create(builder, loc, operands); - yield.erase(); - }; - updateYield(newIfOp.thenYield(), ifYieldOperands); - updateYield(newIfOp.elseYield(), elseYieldOperands); - - int resultIdx = 0; - // Replace old if with the new one. - for (auto result : ifOp.getResults()) { - result.replaceAllUsesWith(newIfOp->getResult(resultIdx++)); - } - // Replace the select with the new return value. - for (arith::SelectOp selectOp : selectOps) { - selectOp.replaceAllUsesWith(newIfOp->getResult(resultIdx++)); - selectOp.erase(); - } - - ifOp.erase(); - } - } -}; - -} // namespace gpu -} // namespace triton -} // namespace mlir diff --git a/third_party/mthreads/lib/Dialect/TritonGPU/Transforms/DecomposeScaledBlocked.cpp b/third_party/mthreads/lib/Dialect/TritonGPU/Transforms/DecomposeScaledBlocked.cpp deleted file mode 100644 index 509b815eb3..0000000000 --- a/third_party/mthreads/lib/Dialect/TritonGPU/Transforms/DecomposeScaledBlocked.cpp +++ /dev/null @@ -1,261 +0,0 @@ -#include "triton/Dialect/TritonGPU/Transforms/DecomposeScaledBlocked.h" - -#include "mlir/IR/Types.h" -#include "mlir/IR/Value.h" -#include "mlir/Support/LogicalResult.h" - -#include "triton/Dialect/Triton/IR/Dialect.h" -#include "triton/Dialect/TritonGPU/IR/Attributes.h" -#include "triton/Dialect/TritonGPU/IR/Dialect.h" -#include "triton/Dialect/TritonGPU/Transforms/Utility.h" - -using namespace mlir; -using namespace mlir::triton; -using namespace mlir::triton::gpu; - -namespace mlir::triton::gpu { - -SmallVector DecomposeScaledBlocked::getTransposeOrder(int rank) { - assert(rank >= 2); - auto transOrder = llvm::to_vector<2>(llvm::seq(rank - 2)); - transOrder.push_back(rank - 1); - transOrder.push_back(rank - 2); - return transOrder; -} - -LogicalResult -DecomposeScaledBlocked::matchAndRewrite(DotScaledOp scaledDotOp, - PatternRewriter &rewriter) const { - if (isa_and_nonnull( - scaledDotOp.getResult().getType().getEncoding())) - return failure(); - - // TODO: add support for m/n packed formats. - if (!scaledDotOp.getLhsKPack() || !scaledDotOp.getRhsKPack()) - return failure(); - // Types - auto computeType = getComputeType(scaledDotOp.getAElemType(), - scaledDotOp.getBElemType(), rewriter); - - auto scaledA = scaleArg(rewriter, scaledDotOp, 0, computeType); - scaledA = cvtDotOperand(rewriter, scaledDotOp, 0, scaledA); - auto scaledB = scaleArg(rewriter, scaledDotOp, 1, computeType); - scaledB = cvtDotOperand(rewriter, scaledDotOp, 1, scaledB); - auto newDot = DotOp::create(rewriter, scaledDotOp.getLoc(), scaledA, scaledB, - scaledDotOp.getC()); - - rewriter.replaceOpWithNewOp(scaledDotOp, - scaledDotOp.getType(), newDot); - return success(); -} - -FloatType -DecomposeScaledBlocked::getComputeType(ScaleDotElemType aType, - ScaleDotElemType bType, - PatternRewriter &rewriter) const { - if (aType == ScaleDotElemType::FP16 || bType == ScaleDotElemType::FP16) - return rewriter.getF16Type(); - return rewriter.getBF16Type(); -} - -TypedValue -DecomposeScaledBlocked::scaleTo16(PatternRewriter &rewriter, - TypedValue scale, - FloatType computeType) const { - auto loc = scale.getLoc(); - auto scaleTy = scale.getType(); - assert(computeType == rewriter.getBF16Type() || - computeType == rewriter.getF16Type()); - - // Choose an fp type that can fit the scale value. - FloatType largeFpType = computeType == rewriter.getF16Type() - ? rewriter.getF32Type() - : computeType; - int intWidth = largeFpType.getIntOrFloatBitWidth(); - auto intType = rewriter.getIntegerType(intWidth); - - auto zexted = - arith::ExtUIOp::create(rewriter, loc, scaleTy.clone(intType), scale); - // getFpMantissaWidth() returns the number of bits in the mantissa plus the - // sign bit! - int shiftValue = largeFpType.getFPMantissaWidth() - 1; - auto shiftConst = - arith::ConstantIntOp::create(rewriter, loc, shiftValue, intWidth); - auto shift = - SplatOp::create(rewriter, loc, scaleTy.clone(intType), shiftConst); - auto shlRes = arith::ShLIOp::create(rewriter, loc, zexted, shift); - Value scaleFP = - BitcastOp::create(rewriter, loc, scaleTy.clone(largeFpType), shlRes); - if (largeFpType != computeType) { - scaleFP = arith::TruncFOp::create(rewriter, loc, scaleTy.clone(computeType), - scaleFP); - } - return cast>(scaleFP); -} - -TypedValue DecomposeScaledBlocked::broadcastScale( - PatternRewriter &rewriter, DotScaledOp scaledDotOp, ModuleOp mod, - TypedValue scale, int dim) const { - auto *ctx = rewriter.getContext(); - auto loc = scale.getLoc(); - auto scaleTy = scale.getType(); - auto rank = scaleTy.getRank(); - // 2.1) Expand dims along the last dimension - { - // 2.1.1) Find default encoding for ExpandDims - auto shape = to_vector(scaleTy.getShape()); - shape.insert(shape.end(), 1); - auto nWarps = lookupNumWarps(scaledDotOp); - auto threadsPerWarp = TritonGPUDialect::getThreadsPerWarp(mod); - auto numCTAs = TritonGPUDialect::getNumCTAs(mod); - auto blockedEnc = - getDefaultBlockedEncoding(ctx, shape, nWarps, threadsPerWarp, numCTAs); - // 2.1.2) Cast scale16 to SliceEncoding - auto sliceEnc = SliceEncodingAttr::get(ctx, rank, blockedEnc); - auto sliceType = scaleTy.cloneWithEncoding(sliceEnc); - scale = ConvertLayoutOp::create(rewriter, loc, sliceType, scale); - } - auto expandScale = ExpandDimsOp::create(rewriter, loc, scale, rank); - // 2.2) Broadcast the dimension to size 32 - auto scaleShape = to_vector(scaleTy.getShape()); - scaleShape.push_back(32); - auto broadcastScale = BroadcastOp::create( - rewriter, loc, expandScale.getType().clone(scaleShape), expandScale); - // 2.3) Transpose the dimension to the scaled dimension - auto transposeOrder = llvm::to_vector(llvm::seq(rank)); - transposeOrder.insert(transposeOrder.begin() + dim + 1, rank); - auto transposedScale = - TransOp::create(rewriter, loc, broadcastScale, transposeOrder); - // 2.4) Reshape to the shape of v - scaleShape.pop_back(); - scaleShape[dim] *= 32; - auto reshapeScale = - ReshapeOp::create(rewriter, loc, scaleShape, transposedScale); - return reshapeScale; -} - -TypedValue DecomposeScaledBlocked::maskNan( - PatternRewriter &rewriter, DotScaledOp scaledDotOp, - TypedValue mxfp, TypedValue scale, - int dim) const { - // Skip NaN checks if fastMath - if (scaledDotOp.getFastMath()) - return mxfp; - - // Implement tl.where(scale == 0xFF, float("nan"), mxfp) - auto loc = scale.getLoc(); - auto mod = scaledDotOp->getParentOfType(); - - // Scale is NaN - auto scaleTy = scale.getType(); - auto constFF = arith::ConstantOp::create( - rewriter, loc, scaleTy, - DenseElementsAttr::get(scaleTy, - APInt(scaleTy.getElementTypeBitWidth(), 0xff))); - auto scaleIsNan = cast>( - arith::CmpIOp::create(rewriter, loc, arith::CmpIPredicate::eq, scale, - constFF) - .getResult()); - auto cond = broadcastScale(rewriter, scaledDotOp, mod, scaleIsNan, dim); - // Make scale is NaN compatible with mxfp - auto condTy = cond.getType(); - condTy = condTy.cloneWithEncoding(mxfp.getType().getEncoding()); - cond = ConvertLayoutOp::create(rewriter, loc, condTy, cond); - - // Create NaN - auto mxfpTy = mxfp.getType(); - auto nan = APFloat::getNaN( - cast(mxfpTy.getElementType()).getFloatSemantics()); - auto constNan = arith::ConstantOp::create( - rewriter, loc, mxfpTy, DenseElementsAttr::get(mxfpTy, nan)); - - auto result = arith::SelectOp::create(rewriter, loc, cond, constNan, mxfp); - return cast>(result.getResult()); -} - -TypedValue -DecomposeScaledBlocked::scaleArg(PatternRewriter &rewriter, - DotScaledOp scaledDotOp, int opIdx, - FloatType computeType) const { - auto v = opIdx == 0 ? scaledDotOp.getA() : scaledDotOp.getB(); - auto scale = opIdx == 0 ? scaledDotOp.getAScale() : scaledDotOp.getBScale(); - auto isFp4 = - ScaleDotElemType::E2M1 == - (opIdx == 0 ? scaledDotOp.getAElemType() : scaledDotOp.getBElemType()); - auto fastMath = scaledDotOp.getFastMath(); - - auto loc = v.getLoc(); - auto rank = v.getType().getRank(); - auto kDim = opIdx == 0 ? rank - 1 : rank - 2; - - // 0) Upcast value to computeType (fp16/bf16) - if (isFp4) { - // We always pack along the fastest moving dimension, kDim - v = Fp4ToFpOp::create(rewriter, loc, v, computeType, kDim); - } else { - auto vType16 = v.getType().clone(computeType); - v = cast>( - FpToFpOp::create(rewriter, loc, vType16, v).getResult()); - } - if (!scale) - return v; - - // 1) Cast scale to fp16/bf16, broadcast it and convert its layout - auto reshapeScale = extendAndBroadcastScale(rewriter, scaledDotOp, scale, - computeType, v.getType(), opIdx); - - // 2) Multiply - auto mxfp = cast>( - arith::MulFOp::create(rewriter, loc, v, reshapeScale).getResult()); - - // 3) If the scale is NaN, return NaN, else return the scaled value. - return maskNan(rewriter, scaledDotOp, mxfp, scale, kDim); -} - -TypedValue DecomposeScaledBlocked::extendAndBroadcastScale( - PatternRewriter &rewriter, DotScaledOp scaledDotOp, - TypedValue &scale, FloatType computeType, - RankedTensorType dstType, int opIdx) const { - auto loc = scale.getLoc(); - auto mod = scaledDotOp->getParentOfType(); - auto v = opIdx == 0 ? scaledDotOp.getA() : scaledDotOp.getB(); - auto rank = v.getType().getRank(); - auto kDim = opIdx == 0 ? rank - 1 : rank - 2; - - // For some weird reason, we take the scale with shape as if it were coming - // from the lhs even when it's the rhs. In a normal world, we should accept - // this parameter transposed, as we do with the mxfp. - // - // Notice: this is an inplace change. - if (opIdx == 1) { - auto order = getTransposeOrder(rank); - scale = TransOp::create(rewriter, loc, scale, order); - } - - // 1) Cast scale to compute type (fp16/bf16) - auto scale16 = scaleTo16(rewriter, scale, computeType); - - // 2) Broadcast scale to the same shape as v and convert the layout - auto reshapeScale = broadcastScale(rewriter, scaledDotOp, mod, scale16, kDim); - return ConvertLayoutOp::create(rewriter, loc, dstType, reshapeScale); -} - -TypedValue -DecomposeScaledBlocked::cvtDotOperand(PatternRewriter &rewriter, - DotScaledOp scaledDotOp, int opIdx, - TypedValue v) const { - auto *ctx = rewriter.getContext(); - auto retEnc = scaledDotOp.getType().getEncoding(); - auto vType = v.getType(); - auto encoding = - DotOperandEncodingAttr::get(ctx, opIdx, retEnc, vType.getElementType()); - auto retTy = vType.cloneWithEncoding(encoding); - return ConvertLayoutOp::create(rewriter, v.getLoc(), retTy, v); -} - -void populateDecomposeScaledBlockedPatterns(RewritePatternSet &patterns, - int benefit) { - patterns.add(patterns.getContext(), benefit); -} - -} // namespace mlir::triton::gpu diff --git a/third_party/mthreads/lib/Dialect/TritonGPU/Transforms/FuseNestedLoops.cpp b/third_party/mthreads/lib/Dialect/TritonGPU/Transforms/FuseNestedLoops.cpp deleted file mode 100644 index 96e3752c6e..0000000000 --- a/third_party/mthreads/lib/Dialect/TritonGPU/Transforms/FuseNestedLoops.cpp +++ /dev/null @@ -1,1222 +0,0 @@ -#include "mlir/Analysis/TopologicalSortUtils.h" -#include "mlir/Dialect/LLVMIR/LLVMDialect.h" -#include "mlir/Dialect/UB/IR/UBOps.h" -#include "mlir/IR/Dominance.h" -#include "mlir/IR/ImplicitLocOpBuilder.h" -#include "mlir/Interfaces/SideEffectInterfaces.h" -#include "mlir/Transforms/LoopInvariantCodeMotionUtils.h" -#include "mlir/Transforms/RegionUtils.h" -#include "triton/Dialect/Triton/IR/Dialect.h" -#include "triton/Dialect/TritonGPU/Transforms/Passes.h" -#include "triton/Dialect/TritonGPU/Transforms/PipeliningUtility.h" -#include "llvm/ADT/STLExtras.h" -#include "llvm/Support/Debug.h" -#include - -namespace mlir { -namespace triton { -namespace gpu { - -//===----------------------------------------------------------------------===// -// Pass Definition -//===----------------------------------------------------------------------===// - -#define GEN_PASS_DEF_TRITONGPUFUSENESTEDLOOPS -#include "triton/Dialect/TritonGPU/Transforms/Passes.h.inc" - -// This attribute is set by the front-end to control whether fusion is on. -static constexpr llvm::StringLiteral kFlattenAttr = "tt.flatten"; -// This attribute indicates the inner loop length has been speculated. -static constexpr llvm::StringLiteral kMustExecuteAttrName = "ttg.must-execute"; -// This attribute is just used for testing the pass. -static constexpr llvm::StringLiteral kAlwaysFuseAttrName = "ttg.always-fuse"; - -namespace { -struct FuseNestedLoopsPass - : public impl::TritonGPUFuseNestedLoopsBase { - using TritonGPUFuseNestedLoopsBase::TritonGPUFuseNestedLoopsBase; - - void runOnOperation() override; -}; - -//===----------------------------------------------------------------------===// -// LoopNest -//===----------------------------------------------------------------------===// - -// A node in the loop nest represents a single for loop with a list of -// immediately nested loops. -struct LoopNestNode { - LoopNestNode(scf::ForOp loop) : loop(loop) {} - - // The for loop. - scf::ForOp loop; - // Loops nested immediately below this loop. - SmallVector children; -}; - -// A loop nest is a tree of loops. -struct LoopNest { - LoopNest(scf::ForOp outermost); - - // Print the loop nest. - void print(raw_ostream &os) const; - // Dump the loop nest for debugging. - LLVM_DUMP_METHOD void dump() const; - - // Owner of the memory of the nodes. - SmallVector> nodes; - - // The outermost loop in the nest, which has no preconditions. Even if the - // outermost loop is contained within an if, its preconditions relative to the - // loop nest are empty. - LoopNestNode *root; -}; -} // namespace - -LoopNest::LoopNest(scf::ForOp outermost) - : root( - nodes.emplace_back(std::make_unique(outermost)).get()) { -} - -void LoopNest::print(raw_ostream &os) const { - // Print just the first line of the loop's textual IR. - std::string buffer; - auto printLoopFirstLine = [&](scf::ForOp loop) { - buffer.clear(); - llvm::raw_string_ostream str(buffer); - loop.print(str); - os << buffer.substr(0, buffer.find('\n')); - }; - - os << "LoopNest:\n"; - SmallVector> stack; - stack.emplace_back(root, 0); - while (!stack.empty()) { - auto [node, indent] = stack.pop_back_val(); - - // Print the current loop. - os << std::string(indent * 2, ' '); - printLoopFirstLine(node->loop); - os << "\n"; - - // Push the children of the current loop. - for (LoopNestNode *child : node->children) - stack.emplace_back(child, indent + 1); - } - os << "\n"; -} - -void LoopNest::dump() const { print(llvm::dbgs()); } - -//===----------------------------------------------------------------------===// -// findLoopNests -//===----------------------------------------------------------------------===// - -// Forward declaration. -static void findLoopNests(Operation *container, - SmallVectorImpl &nests); - -// Recursively construct a loop nest. -static void constructLoopNest(LoopNestNode *parent, LoopNest &nest, - SmallVectorImpl &nests) { - parent->loop->walk([&](Operation *op) { - if (op == parent->loop) - return WalkResult::advance(); - - if (auto forOp = dyn_cast(op)) { - auto &child = - nest.nodes.emplace_back(std::make_unique(forOp)); - parent->children.push_back(child.get()); - // Recurse with the current loop nest. - constructLoopNest(child.get(), nest, nests); - return WalkResult::skip(); - } - - // If the traversal encounters any other operation with regions, restart the - // traversal and construct new loop nests. This means ops like `scf.while` - // divide the analysis domain, but it also means loop fusion won't "see" - // across `scf.if`, for example. - // TODO: Handle loop nests with preconditions. The traversal can keep a - // stack of `scf.if` preconditions while constructing the loop nest. - if (op->getNumRegions()) { - findLoopNests(op, nests); - return WalkResult::skip(); - } - - return WalkResult::advance(); - }); -} - -// Find all the loop nests in the operation. The only region operation that -// allows CFG regions is `tt.func`. That means we can just walk starting from -// the function body and can build loop nests directly off the region trees -// contained in the function -- we don't have to worry about CFGs inside the -// nested region trees. -static void findLoopNests(Operation *container, - SmallVectorImpl &nests) { - container->walk([&](scf::ForOp loop) { - LoopNest nest(loop); - constructLoopNest(nest.root, nest, nests); - nests.push_back(std::move(nest)); - return WalkResult::skip(); - }); -} - -//===----------------------------------------------------------------------===// -// Logue -//===----------------------------------------------------------------------===// - -namespace { -// A prologue or epilogue. -struct Logue { - // Move the ops in the logue before the iterator. - void moveBefore(Block *block, Block::iterator it) { - for (Operation *op : ops) - op->moveBefore(block, it); - } - - // Replace all uses of the logue results with the given values, where `logue` - // comprises all the ops in `containingRegion`. - void replaceAllUsesWith(ValueRange values, Region &containingRegion) { - for (auto [newOut, output] : llvm::zip(values, outputs)) { - // Replace uses of the prologue outputs that are not in the prologue, i.e. - // inside the `then` region where it got spliced. - output.replaceUsesWithIf(newOut, [&](OpOperand &use) { - return !containingRegion.isAncestor(use.getOwner()->getParentRegion()); - }); - } - } - - // Get the number of outputs. - unsigned getNumOutputs() const { return outputs.size(); } - // Get the outputs as a `ValueRange`. - ValueRange getOutputs() const { return outputs; } - // Get the types of the outputs. - TypeRange getOutputTypes() const { return getOutputs().getTypes(); } - - // A contiguous range of ops representing the prologue or epilogue. - SmallVector ops; - // The outputs of the logue. These are the SSA value results of `ops` that are - // used by ops outside of `ops`. - SmallVector outputs; -}; -} // namespace - -// Given a range of ops, form it into a logue by finding the outputs. -static Logue createLogueFrom(llvm::iterator_range ops, - mlir::DominanceInfo &domInfo) { - Logue logue; - for (Operation &op : ops) - logue.ops.push_back(&op); - - if (ops.empty()) - return logue; - - // An op result is an output of the logue if the last operation in the logue - // dominates any of its users. - Operation &lastOp = *std::prev(ops.end()); - auto isOutput = [&](OpResult result) { - for (Operation *user : result.getUsers()) { - if (domInfo.properlyDominates(&lastOp, user)) - return true; - } - return false; - }; - - // Find the outputs. - for (Operation &op : ops) { - for (OpResult result : op.getOpResults()) { - if (isOutput(result)) - logue.outputs.push_back(result); - } - } - - return logue; -} - -//===----------------------------------------------------------------------===// -// fuseOneLevel -//===----------------------------------------------------------------------===// - -// Only hoist operations that are side-effect free and "cheap" (i.e. only scalar -// operands). Importantly, we need to be able to hoist code generated by fusing -// children loops into their parents so the algorithm can be applied -// recursively. This includes integer division, which are not speculatable, but -// we know they will never divide by zero. -static bool canHoistLoopBoundComputation(Operation *op) { - auto isScalar = [](Type type) { - return type.isIntOrIndexOrFloat() || isa(type); - }; - return (isMemoryEffectFree(op) || hasSingleEffect(op)) && - llvm::all_of(op->getOperandTypes(), isScalar) && - llvm::all_of(op->getResultTypes(), isScalar); -} - -// Determine if all of `values` are or can be made invariant to the outer loop -// by hoisting operations. `toHoist` is shared across all child loop bounds. -static bool isOuterLoopInvariant(mlir::DominanceInfo &domInfo, scf::ForOp outer, - ArrayRef values, - llvm::SetVector &toHoist) { - return getDominatingValueSetOpsToHoist( - domInfo, outer, values, toHoist, canHoistLoopBoundComputation, - [&](BlockArgument arg) { - return isa(arg.getOwner()->getParentOp()); - }); -} - -static bool canSliceBounds(mlir::DominanceInfo &domInfo, scf::ForOp outer, - ArrayRef values, - llvm::SetVector &ops) { - return getDominatingValueSetOpsToHoist( - domInfo, outer, values, ops, canHoistLoopBoundComputation, - [&](BlockArgument arg) { - return arg == outer.getInductionVar() || - isa(arg.getOwner()->getParentOp()); - }); -} - -// Pessimistically assume the internal storage bitwidth for index types. -static unsigned getIntTypeWidth(Type type) { - if (isa(type)) - return IndexType::kInternalStorageBitWidth; - return cast(type).getWidth(); -} - -// Generate IR to compute the number of iterations of a loop. -static Value computeNumIters(ImplicitLocOpBuilder &b, Value lowerBound, - Value upperBound, Value step) { - // len(range(lb, ub, step)) = ceildiv(ub - lb, step) - // This works even if step is negative. - Value diff = arith::SubIOp::create(b, upperBound, lowerBound); - // Let someone else prove it can be unsigned. - return arith::CeilDivSIOp::create(b, diff, step); -} - -// Generate IR to compute the number of iterations of a loop. -static Value computeNumIters(ImplicitLocOpBuilder &b, scf::ForOp loop) { - return computeNumIters(b, loop.getLowerBound(), loop.getUpperBound(), - loop.getStep()); -} - -// Cast an integer or index value to an integer or index `type`, if necessary. -static Value castIntIfNecessary(ImplicitLocOpBuilder &b, Value value, - Type type) { - if (value.getType() == type) - return value; - if (isa(value.getType()) || isa(type)) - return arith::IndexCastOp::create(b, type, value); - if (cast(value.getType()).getWidth() > - cast(type).getWidth()) - return arith::TruncIOp::create(b, type, value); - return arith::ExtSIOp::create(b, type, value); -} - -// To model an "undef" value, i.e. a value that is known to never be read on -// live code paths, create a zero-valued constant where possible, otherwise use -// a poison value. PTXAS appears to generate better code with zeros compared to -// poison values. -static Value createPoisonOrZero(ImplicitLocOpBuilder &b, Type type) { - Type elTy = getElementTypeOrSelf(type); - if (!elTy.isIntOrIndexOrFloat() || - (!isa(type) && type != elTy)) - return ub::PoisonOp::create(b, type); - - TypedAttr attr = isa(elTy) ? TypedAttr(b.getFloatAttr(elTy, 0)) - : b.getIntegerAttr(elTy, 0); - if (auto tensor = dyn_cast(type)) - attr = SplatElementsAttr::get(tensor, attr); - return arith::ConstantOp::create(b, attr); -} - -static scf::YieldOp getYield(Region &body) { - return cast(body.front().back()); -} - -static scf::IfOp eraseIfResults(ImplicitLocOpBuilder &b, scf::IfOp ifOp, - llvm::BitVector indices, - SmallVector replaceWith) { - OpBuilder::InsertionGuard guard(b); - b.setInsertionPoint(ifOp); - while (indices.size() < ifOp.getNumResults()) - indices.push_back(false); - - getYield(ifOp.getThenRegion())->eraseOperands(indices); - getYield(ifOp.getElseRegion())->eraseOperands(indices); - - TypeRange newTypes = getYield(ifOp.getThenRegion()).getOperandTypes(); - auto newIf = scf::IfOp::create(b, newTypes, ifOp.getCondition()); - newIf.getThenRegion().takeBody(ifOp.getThenRegion()); - newIf.getElseRegion().takeBody(ifOp.getElseRegion()); - - SmallVector replacements; - auto replIt = replaceWith.begin(); - auto resIt = newIf->result_begin(); - for (unsigned i : llvm::seq(ifOp.getNumResults())) - replacements.push_back(indices[i] ? *replIt++ : *resIt++); - assert(ValueRange(replacements).getTypes() == ifOp.getResultTypes()); - ifOp.replaceAllUsesWith(replacements); - ifOp.erase(); - return newIf; -} - -namespace { -struct InnerLoop { - InnerLoop(scf::ForOp op, llvm::SetVector slicedOps) - : op(op), slicedOps(std::move(slicedOps)) {} - - // Return true if the loop bounds are outer loop invariant. - bool isOuterLoopInvariant() const { return slicedOps.empty(); } - - // The actual loop op. - scf::ForOp op; - // Ops that must be sliced to compute the loop bounds - llvm::SetVector slicedOps; -}; -} // namespace - -// Given a one level loop nest in the form -// -// for i in range(lbi, ubi, stepi): -// prologue0(i) -// for j0 in range(lbj0, ubj0, stepj0): -// body0(i, j0) -// epilogue1(i) -// for j1 in range(lbj1, ubj1, stepj1): -// body1(i, j1) -// epilogue2(i) -// ... -// for jN in range(lbjN, ubjN, stepjN): -// bodyN(i, jN) -// epilogue(i) -// -// Rewrite this into a single loop in the form: -// -// len_i = len(range(lbi, ubi, stepi)) -// len_j0 = len(range(lbj0, ubj0, stepj0)) -// len_j1 = len(range(lbj1, ubj1, stepj1)) -// ... -// len_jN = len(range(lbjN, ubjN, stepjN)) -// inner_len = max(1, len_j0) + max(1, len_j1) + ... + max(1, len_jN) - N -// total_iters = len_i * inner_len -// -// T = 0 -// i = lbi - stepi -// for _ in range(total_iters): -// if T == 0: -// i += stepi -// prologue0(i) -// j0 = lbj0 -// if T >= 0 and T < len_j0: -// body0(i, j0) -// j0 += stepj0 -// -// if T == max(1, len_j0) - 1: -// prologue1(i) -// j1 = lbj1 -// if T >= max(1, len_j0) - 1 -// and T < max(1, len_j0) - 1 + len_j1: -// body1(i, j1) -// j1 += stepj1 -// -// if T == max(1, len_j0) + max(1, len_j1) - 2: -// prologue2(i) -// j2 = lbj2 -// if T >= max(1, len_j0) + max(1, len_j1) - 2 -// and T < max(1, len_j0) + max(1, len_j1) - 2 + len_j2: -// body2(i, j2) -// j2 += stepj2 -// -// ... -// -// if T == max(1, len_j0) + max(1, len_j1) + ... + max(1, len_jN-1) - N: -// prologueN(i) -// jN = lbjN -// if T >= max(1, len_j0) + max(1, len_j1) + ... + max(1, len_jN-1) - N -// and T < max(1, len_j0) + max(1, len_j1) + ... + max(1, len_jN-1) - N + -// len_jN: -// bodyN(i, jN) -// jN += stepjN -// -// if T == max(1, len_j0) + max(1, len_j1) + ... + max(1, len_jN) - (N + 1): -// epilogue(i) -// T = 0 if T == (inner_len - 1) else T + 1 -// -// This routine can be applied recursively on a loop nest tree, leaf-to-root, to -// flatten the loop nest into a single loop. However, this routine only fuses -// child loops whose loop bounds are invariant to the parent loop. For child -// loops where this is not the case, the function will ignore them. -// -// We could fuse loops with parent-loop-variant or even data-dependent bounds, -// but this will require generating `scf.while` in a form that is not friendly -// to the pipeliner. In order to effectively fuse and pipeline these kinds of -// loop nests, loop nest fusion and the pipeliner need to share a higher-level -// representation (or perhaps be the same pass). -// -// Note that there are many potential forms of the fused loop. This routine will -// attempt to minimize the number of fused loop iterations by overlapping the -// iteration spaces of the child loops and the epilogues. E.g. the last -// iteration of bodyjK will execute on the same fused loop iteration as -// epilogueK and the first iteration of bodyj(K+1). Hence the `- N` term in the -// total number of iterations. -// -// What the above Python-pseudo-code glosses over is SSA dependency management. -// To interpret the pseudocode as SSA IR, just imagine everything is put back -// into allocas and SSA formation re-runs after fusion, which one should note -// will introduce undefs. -// -// Handling dependencies will require turning implicit captures into -// loop-carried dependencies. Consider: -// -// scf.for %i = %lbi to %ubi step %stepi { -// %a = tt.call @func(%i) -// scf.for %j = %lbj to %ubj step %stepj { -// %b = tt.call @use(%a, %j) -// } -// } -// -// This needs to be rewritten into: -// -// %poison = ub.poison -// %Tlast, %ilast, %jlast, %alast = scf.for %unused = ... -// iter_args(%Tprev = %c-1_i32, -// %iprev = %lbi - %stepi, -// %jprev = %poison, -// %aprev = %poison) -> (i32, i32, i32, i32) { -// %T = (%Tprev + 1) mod (...) -// %a, %i, %j = scf.if %T == 0 { -// %inext = %iprev + 1 -// %jnext = %lbj - %stepj -// -// %anext = tt.call @func(%i) -// yield %inext, %jnext, %anext -// } else { -// yield %iprev, %jprev, %aprev -// } -// -// scf.if %T >= 0 and %T < ... { -// tt.call @use(%a, %j) -// } -// -// Note: the induction variables will be initialized to their lower bound to -// avoid underflow in lbjk - stepjk, with the exception of the outer loop -// induction variable, which needs to be incremented inside the prologue to -// avoid a dependency on the epilogue. This helps the scheduler behave. -// -// Any inputs and outputs of the loop bodies would also need to be handled -// similarly: initialized as undef if appropriate and carried through the fused -// loop. This is why fusion will increase liveranges. To minimize the number of -// additional loop-carried values, the routine will analyze the subblock of IR -// inside each `prologueK` and determine its "outputs" as intermediate SSA -// values that are used later in the loop nest. -static void fuseOneLevel(LoopNestNode *parent, mlir::DominanceInfo &domInfo) { - scf::ForOp outer = parent->loop; - - SmallVector innerLoops; - llvm::SetVector toHoist; - for (LoopNestNode *child : parent->children) { - scf::ForOp inner = child->loop; - assert(child->children.empty() && "fuseOneLevel runs leaf-to-root"); - - // Check if the inner loop bounds are or can be made invariant to the outer - // loop. Check them all at once to avoid adding ops to `toHoist` if not - // necessary. - if (isOuterLoopInvariant( - domInfo, outer, - {inner.getLowerBound(), inner.getUpperBound(), inner.getStep()}, - toHoist)) { - // Add this child to the list of loops to fuse. - innerLoops.push_back({child->loop, {}}); - continue; - } - - // Check if the loop bounds can be sliced. - llvm::SetVector slicedOps; - if (canSliceBounds( - domInfo, outer, - {inner.getLowerBound(), inner.getUpperBound(), inner.getStep()}, - slicedOps)) { - innerLoops.push_back({child->loop, std::move(slicedOps)}); - continue; - } - } - - // From the perspective of the overall analysis, we can delete all the - // children of the current loop node. Child loops that cannot be fused are now - // treated opaquely by the rest of the analysis. This allows partial fusing of - // the constructed loop nest. - parent->children.clear(); - - // If there are no child loops to fuse, then there is nothing to do. - if (innerLoops.empty()) - return; - - // The transformation will definitely succeed on `childrenToFuse`. `toHoist` - // only contains the operations that must be hoisted for `childrenToFuse` to - // be fusible. - hoistOpsBefore(outer, toHoist); - - // Determine the integer type to use for the length computations. Use an - // integer bitwidth twice the size of the largest integer, up to 64 bits, to - // avoid overflow. - unsigned intTyWidth = getIntTypeWidth(outer.getInductionVar().getType()); - - // Generate the computations of the fused loop bounds. - Location loc = outer.getLoc(); - ImplicitLocOpBuilder b(loc, outer); - for (InnerLoop &loop : innerLoops) { - intTyWidth = std::max(intTyWidth, - getIntTypeWidth(loop.op.getInductionVar().getType())); - } - auto intTy = b.getIntegerType(intTyWidth); - bool allInvariant = llvm::all_of( - innerLoops, [](InnerLoop &loop) { return loop.isOuterLoopInvariant(); }); - - Value lenOuter = computeNumIters(b, outer); - SmallVector lenInners; - for (InnerLoop &loop : innerLoops) { - // len_jk = len(range(lbjk, ubjk, stepjk)) - Value lenInner; - if (loop.isOuterLoopInvariant()) - lenInner = castIntIfNecessary(b, computeNumIters(b, loop.op), intTy); - else - lenInner = createPoisonOrZero(b, intTy); - lenInners.push_back(lenInner); - } - - auto intTyCst = [&](int64_t v) { - return arith::ConstantOp::create(b, IntegerAttr::get(intTy, v)); - }; - - // inner_len = max(1, len_j0) + max(1, len_j1) + ... + max(1, len_jN) - N - unsigned N = innerLoops.size() - 1; - Value innerLen = intTyCst(0); - for (auto [loop, lenInner] : llvm::zip(innerLoops, lenInners)) { - if (!loop.isOuterLoopInvariant()) - continue; - innerLen = arith::AddIOp::create( - b, innerLen, arith::MaxSIOp::create(b, intTyCst(1), lenInner)); - } - innerLen = arith::SubIOp::create(b, innerLen, intTyCst(N)); - - // total_iters = len_i * inner_len - Value totalIters = arith::MulIOp::create( - b, castIntIfNecessary(b, lenOuter, intTy), innerLen); - - // Generate a loop to compute the total number of iterations for inner loops - // whose bounds are not outer loop invariant. - IRMapping mapping; - auto peeledLen = - scf::ForOp::create(b, outer.getLowerBound(), outer.getUpperBound(), - outer.getStep(), {totalIters}); - totalIters = peeledLen.getRegionIterArg(0); - mapping.map(outer.getInductionVar(), peeledLen.getInductionVar()); - b.setInsertionPointToStart(peeledLen.getBody()); - for (InnerLoop &loop : innerLoops) { - if (loop.isOuterLoopInvariant()) - continue; - // Cloned the sliced ops into the peeled loop. - for (Operation *op : topologicalSort(loop.slicedOps)) { - if (!mapping.contains(op)) - b.clone(*op, mapping); - } - Value numIters = - computeNumIters(b, mapping.lookupOrDefault(loop.op.getLowerBound()), - mapping.lookupOrDefault(loop.op.getUpperBound()), - mapping.lookupOrDefault(loop.op.getStep())); - numIters = castIntIfNecessary(b, numIters, intTy); - // Accumulate into the total number of iterations. - numIters = arith::MaxSIOp::create(b, intTyCst(1), numIters); - totalIters = arith::AddIOp::create(b, totalIters, numIters); - } - scf::YieldOp::create(b, totalIters); - totalIters = peeledLen.getResults().front(); - b.setInsertionPointAfter(peeledLen); - - // The outputs of the prologue, each epilogue, and all inner loop bodies need - // to carried through the fused loop. - SmallVector logues; - auto addLogue = [&](Block::iterator begin, Block::iterator end) { - logues.push_back(createLogueFrom({begin, end}, domInfo)); - }; - // prologue0 - addLogue(outer.getBody()->begin(), innerLoops.front().op->getIterator()); - // prologuek where 0 < k <= N - for (auto i : llvm::seq(0, innerLoops.size() - 1)) { - addLogue(std::next(innerLoops[i].op->getIterator()), - innerLoops[i + 1].op->getIterator()); - } - // epilogue - addLogue(std::next(innerLoops.back().op->getIterator()), - // Don't include the outer loop yield. - std::prev(outer.getBody()->end())); - - // We need iter args for: - // - The fused loop induction var - // - The outer loop induction var - // - The outer loop iter args - // - The induction vars for each inner loop - // - The outputs of each child loop - // - The outputs of each logue - SmallVector fusedInits; - - // T = 0 - fusedInits.push_back(intTyCst(0)); - // i = lbi - stepi - fusedInits.push_back( - arith::SubIOp::create(b, outer.getLowerBound(), outer.getStep())); - - unsigned outerArgsStartIdx = fusedInits.size(); - llvm::append_range(fusedInits, outer.getInits()); - unsigned lenInnersStartIdx = fusedInits.size(); - llvm::append_range(fusedInits, lenInners); - unsigned innerLenStartIdx = fusedInits.size(); - fusedInits.push_back(innerLen); - - // Everything else is initialized to undef. - unsigned ivarStartIdx = fusedInits.size(); - for (InnerLoop &loop : innerLoops) { - fusedInits.push_back( - createPoisonOrZero(b, loop.op.getInductionVar().getType())); - } - unsigned innerOutsStartIdx = fusedInits.size(); - for (InnerLoop &loop : innerLoops) { - for (Type resultType : loop.op.getResultTypes()) - fusedInits.push_back(createPoisonOrZero(b, resultType)); - } - unsigned logueOutsStartIdx = fusedInits.size(); - for (Logue &logue : llvm::drop_end(logues)) { - for (Type outputType : logue.getOutputTypes()) - fusedInits.push_back(createPoisonOrZero(b, outputType)); - } - - // for _ in range(total_iters): - auto fused = - scf::ForOp::create(b, intTyCst(0), totalIters, intTyCst(1), fusedInits); - // Replace the outer loop args with the args in the fused loop args. - for (auto [arg, fusedArg] : - llvm::zip(outer.getRegionIterArgs(), - fused.getRegionIterArgs().slice(outerArgsStartIdx))) { - arg.replaceAllUsesWith(fusedArg); - } - ValueRange lenInnersRange = - fused.getRegionIterArgs().slice(lenInnersStartIdx, lenInners.size()); - for (auto [lenInner, lenInnerArg] : llvm::zip(lenInners, lenInnersRange)) - lenInner = lenInnerArg; - b.setInsertionPointToStart(fused.getBody()); - - Value T = fused.getRegionIterArg(0); - // `i` is computed inside the first prologue. - Value curI = fused.getRegionIterArg(1); - Value i; - - auto lenInnersIt = - ValueRange(fused.getRegionIterArgs()).begin() + lenInnersStartIdx; - - ArrayRef ivars = fused.getRegionIterArgs().slice(ivarStartIdx); - auto bodyOutsIt = - ValueRange(fused.getRegionIterArgs()).begin() + innerOutsStartIdx; - auto logueOutsIt = - ValueRange(fused.getRegionIterArgs()).begin() + logueOutsStartIdx; - SmallVector prologueIfs, bodyIfs; - for (unsigned k = 0; k <= N; ++k) { - // if T == max(1, len_j0) + ... max(1, len_jk-1) - k - // [[if k == 0]] i += stepi - // prologuek(i) - // jk = lbjk - Value innerStartT = intTyCst(0); - for (unsigned i = 0; i < k; ++i) { - innerStartT = arith::AddIOp::create( - b, innerStartT, arith::MaxSIOp::create(b, intTyCst(1), lenInners[i])); - } - innerStartT = arith::SubIOp::create(b, innerStartT, intTyCst(k)); - Value prologueCond = - arith::CmpIOp::create(b, arith::CmpIPredicate::eq, T, innerStartT); - - // The `scf.if` outputs will be `jk` and the outputs of prologuek. We also - // have to initialize the inner loop iter args. - scf::ForOp inner = innerLoops[k].op; - Logue &prologue = logues[k]; - - SmallVector prologueOutTypes{inner.getInductionVar().getType()}; - llvm::append_range(prologueOutTypes, prologue.getOutputTypes()); - llvm::append_range(prologueOutTypes, inner.getInits().getTypes()); - if (k == 0) { - prologueOutTypes.push_back(curI.getType()); - prologueOutTypes.append(innerLoops.size(), intTy); - prologueOutTypes.push_back(innerLen.getType()); - } - auto prologueIf = scf::IfOp::create(b, prologueOutTypes, prologueCond); - prologueIfs.push_back(prologueIf); - - // Splice prologuek into the `then` region. - Block *thenBlock = b.createBlock(&prologueIf.getThenRegion()); - prologue.moveBefore(thenBlock, thenBlock->end()); - - if (k == 0) { - // Increment `i` and replace its uses inside the prologue. - b.setInsertionPointToStart(thenBlock); - i = arith::AddIOp::create(b, curI, outer.getStep()); - mlir::replaceAllUsesInRegionWith(outer.getInductionVar(), i, - prologueIf.getThenRegion()); - - // Compute the variant inner loop lengths. - IRMapping mapping; - for (auto [loop, lenInner] : llvm::zip(innerLoops, lenInners)) { - if (loop.isOuterLoopInvariant()) - continue; - for (Operation *op : topologicalSort(loop.slicedOps)) { - if (!mapping.contains(op)) - b.clone(*op, mapping); - } - lenInner = - computeNumIters(b, mapping.lookupOrDefault(loop.op.getLowerBound()), - mapping.lookupOrDefault(loop.op.getUpperBound()), - mapping.lookupOrDefault(loop.op.getStep())); - lenInner = castIntIfNecessary(b, lenInner, intTy); - innerLen = arith::AddIOp::create( - b, innerLen, arith::MaxSIOp::create(b, intTyCst(1), lenInner)); - } - } - - // Yield the initialized jk, the prologue outputs, and the initial values of - // the inner loop. - b.setInsertionPointToEnd(thenBlock); - SmallVector thenOuts{inner.getLowerBound()}; - llvm::append_range(thenOuts, prologue.getOutputs()); - llvm::append_range(thenOuts, inner.getInits()); - if (k == 0) { - thenOuts.push_back(i); - llvm::append_range(thenOuts, lenInners); - thenOuts.push_back(innerLen); - } - scf::YieldOp::create(b, thenOuts); - - // In the `else` region, just yield the last values of jk, the outputs, and - // the iter args. - b.createBlock(&prologueIf.getElseRegion()); - Value lastJk = ivars[k]; - unsigned numOuts = prologue.getNumOutputs(); - SmallVector elseOuts{lastJk}; - elseOuts.append(logueOutsIt, logueOutsIt + numOuts); - elseOuts.append(bodyOutsIt, bodyOutsIt + inner.getNumResults()); - if (k == 0) { - elseOuts.push_back(curI); - llvm::append_range(elseOuts, lenInnersRange); - // Peephole the passthrough of `innerLen` since MLIR will not optimize it - // away for us. - elseOuts.push_back( - allInvariant ? innerLen : fused.getRegionIterArg(innerLenStartIdx)); - } - logueOutsIt += numOuts; - scf::YieldOp::create(b, elseOuts); - - // The results of the `scf.if` become the values of jk and the prologue - // outputs for the rest of the fused loop. - Value jk = prologueIf.getResult(0); - ValueRange prologueOuts = prologueIf.getResults().slice(1, numOuts); - ValueRange prologueInits = - prologueIf.getResults().slice(1 + numOuts, inner.getNumResults()); - inner.getInductionVar().replaceAllUsesWith(jk); - prologue.replaceAllUsesWith(prologueOuts, prologueIf.getThenRegion()); - for (auto [init, iterArg] : - llvm::zip(prologueInits, inner.getRegionIterArgs())) - iterArg.replaceAllUsesWith(init); - // Replace uses of `i` elsewhere with the prologue result. - if (k == 0) { - ValueRange results = prologueIf.getResults(); - i = results.drop_back(1 + lenInners.size()).back(); - lenInners = results.drop_back().take_back(lenInners.size()); - innerLen = results.back(); - outer.getInductionVar().replaceAllUsesWith(i); - } - - // if T >= max(1, len_j0) + max(1, len_j1) + ... + max(1, len_jk-1) - k - // and T < max(1, len_j0) + max(1, len_j1) + ... + max(1, len_jk-1) - k + - // len_jk - // bodyk(i, jk) - // jk += stepjk - b.setInsertionPointAfter(prologueIf); - Value innerEndT = arith::AddIOp::create( - b, innerStartT, castIntIfNecessary(b, lenInners[k], intTy)); - Value ge = - arith::CmpIOp::create(b, arith::CmpIPredicate::sge, T, innerStartT); - Value lt = - arith::CmpIOp::create(b, arith::CmpIPredicate::slt, T, innerEndT); - Value bodyCond = arith::AndIOp::create(b, ge, lt); - - // The outputs will be the outputs of the inner loop body and the next jk. - SmallVector bodyOutTypes{jk.getType()}; - llvm::append_range(bodyOutTypes, inner->getResultTypes()); - auto bodyIf = scf::IfOp::create(b, bodyOutTypes, bodyCond); - bodyIfs.push_back(bodyIf); - - // Splice bodyk into the `then` region. - inner.getBody()->eraseArguments([](Value arg) { return true; }); - bodyIf.getThenRegion().takeBody(inner.getBodyRegion()); - auto yield = getYield(bodyIf.getThenRegion()); - b.setInsertionPoint(yield); - Value nextJk = arith::AddIOp::create(b, jk, inner.getStep()); - yield->insertOperands(0, nextJk); - - // The `else` region just forwards the values. - b.createBlock(&bodyIf.getElseRegion()); - SmallVector bodyForwardedOuts{jk}; - bodyForwardedOuts.append(bodyOutsIt, bodyOutsIt + inner.getNumResults()); - bodyOutsIt += inner->getNumResults(); - scf::YieldOp::create(b, bodyForwardedOuts); - - // Now we can replace the results of the inner loop with the outputs of the - // body if. - inner.replaceAllUsesWith( - bodyIf.getResults().slice(1, inner.getNumResults())); - - // If the inner loop must execute, then its body does not have to be wrapped - // in a conditional. - if (inner->hasAttr(kMustExecuteAttrName)) { - b.setInsertionPoint(bodyIf); - bodyIf.getConditionMutable().assign( - arith::ConstantOp::create(b, b.getBoolAttr(true))); - } - - // Move the insertion point for the next iteration. - b.setInsertionPointAfter(bodyIf); - } - - // if T == len_j0 + len_j1 + ... + len_jN - N - 1: - // epilogue(i) - Logue &epilogue = logues.back(); - - // The only possible use of an epilogue output is the yield. - auto outerYield = cast(outer.getBody()->getTerminator()); - SmallVector usedIterArgs; - for (Value output : epilogue.getOutputs()) { - for (OpOperand &use : output.getUses()) { - if (use.getOwner() == outerYield) { - usedIterArgs.push_back(fused.getRegionIterArgs().drop_front( - outerArgsStartIdx)[use.getOperandNumber()]); - } - } - } - - auto epilogueCond = - arith::CmpIOp::create(b, arith::CmpIPredicate::eq, T, - arith::SubIOp::create(b, innerLen, intTyCst(1))); - auto epilogueIf = - scf::IfOp::create(b, epilogue.getOutputTypes(), epilogueCond); - - Block *thenBlock = b.createBlock(&epilogueIf.getThenRegion()); - epilogue.moveBefore(thenBlock, thenBlock->end()); - - b.setInsertionPointToEnd(thenBlock); - scf::YieldOp::create(b, epilogue.getOutputs()); - b.createBlock(&epilogueIf.getElseRegion()); - scf::YieldOp::create(b, usedIterArgs); - epilogue.replaceAllUsesWith(epilogueIf.getResults(), - epilogueIf.getThenRegion()); - - // T = 0 if T == (inner_len - 1) else T + 1 - b.setInsertionPointToEnd(fused.getBody()); - Value nextT = arith::AddIOp::create(b, T, intTyCst(1)); - Value rollover = - arith::CmpIOp::create(b, arith::CmpIPredicate::eq, T, - arith::SubIOp::create(b, innerLen, intTyCst(1))); - T = arith::SelectOp::create(b, rollover, intTyCst(0), nextT); - - // Finally, create the yield of the fused loop. - SmallVector outerOuts{T, i}; - llvm::append_range(outerOuts, outerYield.getOperands()); - llvm::append_range(outerOuts, lenInners); - outerOuts.push_back(innerLen); - for (scf::IfOp bodyIf : bodyIfs) - outerOuts.push_back(/*jk=*/bodyIf.getResult(0)); - for (auto [bodyIf, loop] : llvm::zip(bodyIfs, innerLoops)) { - llvm::append_range(outerOuts, - bodyIf.getResults().slice(1, loop.op.getNumResults())); - } - for (auto [logueIf, logue] : llvm::zip(prologueIfs, llvm::drop_end(logues))) { - llvm::append_range(outerOuts, - logueIf.getResults().slice(1, logue.getNumOutputs())); - } - - scf::YieldOp::create(b, outerOuts); - outer.replaceAllUsesWith( - fused.getResults().slice(outerArgsStartIdx, outer.getNumResults())); - - // Reduce dependencies across inner loops by hoisting the initialization of - // inner loop iter args to the outer loop when possible, and then placing the - // reset of these values in the epilogue. - auto fusedInitsIt = fused.getInitsMutable().begin() + innerOutsStartIdx; - auto fusedArgsIt = fused.getRegionIterArgs().begin() + innerOutsStartIdx; - auto fusedYieldIt = getYield(fused.getBodyRegion())->getOpOperands().begin() + - innerOutsStartIdx; - SmallVector yieldsToUpdate; - SmallVector reset, forwarded; - for (auto [loop, ifOp, bodyIf, prologue] : - llvm::zip(innerLoops, prologueIfs, bodyIfs, logues)) { - unsigned numResults = loop.op.getNumResults(); - unsigned prologueSkip = 1 + prologue.getNumOutputs(); - - llvm::BitVector removeIndices(prologueSkip + numResults); - SmallVector replaceWith; - for (auto [i, init] : llvm::enumerate(loop.op.getInits())) { - if (init.getParentRegion() == &fused.getBodyRegion()) - continue; - // Initialize this in the outer loop. - fusedInitsIt[i].assign(init); - replaceWith.push_back(fusedArgsIt[i]); - removeIndices.set(prologueSkip + i); - yieldsToUpdate.push_back(&fusedYieldIt[i]); - forwarded.push_back(bodyIf.getResult(1 + i)); - reset.push_back(init); - } - // Remove the initializers in the corresponding prologue. - eraseIfResults(b, ifOp, removeIndices, replaceWith); - - fusedInitsIt += numResults; - fusedArgsIt += numResults; - fusedYieldIt += numResults; - } - if (!yieldsToUpdate.empty()) { - MutableOperandRange(getYield(epilogueIf.getThenRegion())).append(reset); - MutableOperandRange(getYield(epilogueIf.getElseRegion())).append(forwarded); - b.setInsertionPoint(epilogueIf); - TypeRange newTypes = getYield(epilogueIf.getThenRegion()).getOperandTypes(); - auto newIf = scf::IfOp::create(b, newTypes, epilogueIf.getCondition()); - newIf.getThenRegion().takeBody(epilogueIf.getThenRegion()); - newIf.getElseRegion().takeBody(epilogueIf.getElseRegion()); - epilogueIf.replaceAllUsesWith( - newIf.getResults().take_front(epilogueIf.getNumResults())); - ResultRange newResults = - newIf.getResults().drop_front(epilogueIf.getNumResults()); - for (auto [i, yieldOperand] : llvm::enumerate(yieldsToUpdate)) - yieldOperand->set(newResults[i]); - epilogueIf.erase(); - } - - // Propagate warp specialization flags. - if (outer->hasAttr(kWarpSpecializeAttrName) || - llvm::any_of(innerLoops, [](InnerLoop &loop) { - return loop.op->hasAttr(kWarpSpecializeAttrName); - })) - fused->setAttr(kWarpSpecializeAttrName, b.getUnitAttr()); - - // Propagate the `tt.disallow_acc_multi_buffer` attribute to the parent loop. - bool disallowAccMultiBuffer = getDisallowAccMultiBuffer(outer); - for (InnerLoop &loop : innerLoops) { - disallowAccMultiBuffer |= getDisallowAccMultiBuffer(loop.op); - } - if (disallowAccMultiBuffer) - fused->setAttr(kDisallowAccMultiBufferAttrName, b.getUnitAttr()); - - // Update the parent's loop to the fused loop. Set the new stage count to the - // max stage count of the inner loops. - int numStages = 1; - if (auto stageAttr = outer->getAttrOfType(kNumStagesAttrName)) - numStages = stageAttr.getInt(); - for (InnerLoop &loop : innerLoops) { - if (auto stageAttr = - loop.op->getAttrOfType(kNumStagesAttrName)) - numStages = std::max(numStages, stageAttr.getInt()); - loop.op.erase(); - } - outer.erase(); - parent->loop = fused; - if (numStages > 1) - fused->setAttr(kNumStagesAttrName, b.getI32IntegerAttr(numStages)); -} - -//===----------------------------------------------------------------------===// -// flattenLoopNest -//===----------------------------------------------------------------------===// - -// Completely flatten a loop nest by recursively fusing loops in a post-order -// traversal with `fuseOneLevel`. -static void flattenLoopNest(LoopNestNode *node, mlir::DominanceInfo &domInfo) { - for (LoopNestNode *child : node->children) - flattenLoopNest(child, domInfo); - fuseOneLevel(node, domInfo); -} - -//===----------------------------------------------------------------------===// -// Pass Implementation -//===----------------------------------------------------------------------===// - -// Fuse simple loop nests with a single outer and inner loop, and where the -// inner loop has a `tt.dot` operation. -static bool shouldFuse(const LoopNest &nest) { - if (nest.root->loop->hasAttr(kAlwaysFuseAttrName)) - return true; - - // Only fuse simple loop nests. - return nest.nodes.size() == 2 && nest.root->children.size() == 1 && - nest.root->loop->hasAttr(kFlattenAttr); -} - -// This function identifies a subgraph of cheap ops that can be sunk between two -// regions in the loop nest and moves them, reducing their liveranges. -static void sinkOps(Region &limit, Block *sinkBlock, Block::iterator sinkBefore, - llvm::iterator_range prologue, - function_ref inSinkRegion) { - llvm::SetVector sunkOps; - auto canBeSunk = [&](Operation &op) -> std::pair { - if (!isPure(&op) || isa(op)) - return {false, false}; - // An op can be sunk if all its users are inside the inner loop or are - // marked for sinking. - bool isRoot = true; - for (Operation *user : op.getUsers()) { - if (inSinkRegion(user)) - continue; - isRoot = false; - if (sunkOps.contains(user)) - continue; - return {false, false}; - } - return {true, isRoot}; - }; - - // Find the subgraph of operations that can be sunk. - SmallVector roots; - for (Operation &op : llvm::reverse(prologue)) { - auto [canSink, isRoot] = canBeSunk(op); - if (canSink) - sunkOps.insert(&op); - if (isRoot) - roots.push_back(&op); - } - if (sunkOps.empty()) - return; - - hoistOpsBefore(sinkBlock, sinkBefore, sunkOps); -} - -// Sink ops from the prologue into the epilogue when possible. -static void optimizeEpilogueDependencies(scf::ForOp outerLoop, - scf::ForOp innerLoop, - mlir::DominanceInfo &domInfo) { - auto inEpilogue = [&](Operation *op) { - return domInfo.properlyDominates(innerLoop, op, /*enclosingOpOk=*/false); - }; - Region &limit = outerLoop.getBodyRegion(); - sinkOps(limit, outerLoop.getBody(), std::next(innerLoop->getIterator()), - {outerLoop.getBody()->begin(), innerLoop->getIterator()}, inEpilogue); -} - -// Crudely match llvm.assume(ub > lb) or llvm.assume(lb < ub). -static LogicalResult matchPositiveTripCount(scf::ForOp loop) { - for (Operation *user : loop.getUpperBound().getUsers()) { - if (auto cmp = dyn_cast(user)) { - if (llvm::none_of(cmp->getUsers(), - [](Operation *op) { return isa(op); })) - continue; - if (cmp.getPredicate() == (loop.getUnsignedCmp() - ? arith::CmpIPredicate::ugt - : arith::CmpIPredicate::sgt) && - cmp.getLhs() == loop.getUpperBound() && - cmp.getRhs() == loop.getLowerBound()) - return success(); - if (cmp.getPredicate() == (loop.getUnsignedCmp() - ? arith::CmpIPredicate::ult - : arith::CmpIPredicate::slt) && - cmp.getLhs() == loop.getLowerBound() && - cmp.getRhs() == loop.getUpperBound()) - return success(); - } - } - return failure(); -} - -// Speculate the length of the inner loop such that the loop is known to execute -// at least once. This way, the inner loop body does not have to be placed -// inside a conditional in the fused loop, which interacts better with the -// pipeliner. -static LogicalResult speculateInnerLoopLength(scf::ForOp outerLoop, - scf::ForOp innerLoop, - mlir::DominanceInfo &domInfo) { - Location loc = innerLoop.getLoc(); - ImplicitLocOpBuilder b(loc, outerLoop); - - // Check if the inner loop is known to execute at least once. - if (succeeded(matchPositiveTripCount(innerLoop))) { - innerLoop->setAttr(kMustExecuteAttrName, b.getUnitAttr()); - return success(); - } - - // The inner loop bounds must be outer-loop invariant to speculate from - // outside the loop nest. - llvm::SetVector toHoist; - if (!isOuterLoopInvariant(domInfo, outerLoop, - {innerLoop.getLowerBound(), - innerLoop.getUpperBound(), innerLoop.getStep()}, - toHoist)) - return failure(); - - // Hoist the inner loop bounds computations if necessary. - hoistOpsBefore(outerLoop, toHoist); - - // Mark the inner loop. - innerLoop->setAttr(kMustExecuteAttrName, b.getUnitAttr()); - - // Speculate on whether the length of the inner loop is zero. - Value lenInner = computeNumIters(b, innerLoop); - auto zeroAttr = IntegerAttr::get(lenInner.getType(), 0); - Value innerLoopEmpty = - arith::CmpIOp::create(b, arith::CmpIPredicate::eq, lenInner, - arith::ConstantOp::create(b, zeroAttr)); - auto ifOp = scf::IfOp::create(b, outerLoop.getResultTypes(), innerLoopEmpty); - - // In the `then` branch, the inner loop does not execute. Clone the loop nest - // into it and remove the inner loop. - mlir::IRMapping map; - b.createBlock(&ifOp.getThenRegion()); - auto newLoop = cast(b.clone(*outerLoop, map)); - scf::YieldOp::create(b, newLoop.getResults()); - auto newInnerLoop = cast(map.lookup(innerLoop)); - newInnerLoop.replaceAllUsesWith(newInnerLoop.getInits()); - newInnerLoop.erase(); - - // Clear up the warp specialization attributes for the specialized loop. - newLoop->removeAttr(kWarpSpecializeAttrName); - - // Move the loop nest into the `else` branch. - outerLoop.replaceAllUsesWith(ifOp.getResults()); - Block *block = b.createBlock(&ifOp.getElseRegion()); - outerLoop->remove(); - b.insert(outerLoop); - scf::YieldOp::create(b, outerLoop.getResults()); - - return success(); -} - -static LogicalResult preprocessLoopNest(const LoopNest &nest, - mlir::DominanceInfo &domInfo) { - assert(nest.nodes.size() == 2 && nest.root->children.size() == 1); - - scf::ForOp &outerLoop = nest.root->loop; - scf::ForOp &innerLoop = nest.root->children.front()->loop; - - moveLoopInvariantCode(outerLoop); - optimizeEpilogueDependencies(outerLoop, innerLoop, domInfo); - return speculateInnerLoopLength(outerLoop, innerLoop, domInfo); -} - -void FuseNestedLoopsPass::runOnOperation() { - auto &domInfo = getAnalysis(); - - for (auto func : getOperation().getOps()) { - SmallVector nests; - findLoopNests(func, nests); - for (LoopNest &nest : nests) { - if (!shouldFuse(nest)) - continue; - if (!nest.root->loop->hasAttr(kAlwaysFuseAttrName) && - failed(preprocessLoopNest(nest, domInfo))) - continue; - flattenLoopNest(nest.root, domInfo); - } - } -} - -} // namespace gpu -} // namespace triton -} // namespace mlir diff --git a/third_party/mthreads/lib/Dialect/TritonGPU/Transforms/HoistTMEMAlloc.cpp b/third_party/mthreads/lib/Dialect/TritonGPU/Transforms/HoistTMEMAlloc.cpp deleted file mode 100644 index 86e5e2e774..0000000000 --- a/third_party/mthreads/lib/Dialect/TritonGPU/Transforms/HoistTMEMAlloc.cpp +++ /dev/null @@ -1,586 +0,0 @@ -#include "mlir/IR/Dominance.h" -#include "mlir/Transforms/GreedyPatternRewriteDriver.h" -#include "triton/Dialect/TritonGPU/IR/Dialect.h" -#include "triton/Dialect/TritonGPU/Transforms/MMAv5PipelineUtility.h" -#include "triton/Dialect/TritonGPU/Transforms/Passes.h" -#include "triton/Dialect/TritonGPU/Transforms/PipeliningUtility.h" -#include "triton/Dialect/TritonGPU/Transforms/Utility.h" -#include "triton/Dialect/TritonNvidiaGPU/IR/Dialect.h" - -using namespace mlir; -namespace tt = mlir::triton; -namespace ttg = mlir::triton::gpu; -namespace ttng = mlir::triton::nvidia_gpu; - -namespace mlir { -namespace triton { -namespace gpu { - -#define GEN_PASS_DEF_TRITONGPUHOISTTMEMALLOC -#include "triton/Dialect/TritonGPU/Transforms/Passes.h.inc" - -namespace { - -// This CRTP class is an operation type constraint that checks that it has TMEM -// dependency tokens present. HoistTMEMAlloc requires that TMEM tokens are -// present to check aliasing for its transformations. -template struct HasToken : public OpT { - using OpT::OpT; - - static bool classof(Operation *op) { - if (auto tmemOp = dyn_cast(op)) - return !!tmemOp.getToken(); - return false; - } -}; - -using TMEMTokenLoadOp = HasToken; -using TMEMTokenStoreOp = HasToken; -using TMEMTokenAllocOp = HasToken; - -class CombineTMEMStoreAndSelect : public OpRewritePattern { -public: - using OpRewritePattern::OpRewritePattern; - - LogicalResult matchAndRewrite(ttng::TMEMStoreOp store, - PatternRewriter &rewriter) const override { - if (!store.getDep()) - return failure(); - Value src = store.getSrc(); - auto select = src.getDefiningOp(); - if (!select) { - return failure(); - } - enum { kTrue, kFalse, kUnknown } valueFromTMEM = kUnknown; - Value trueSrc = select.getTrueValue(); - Value falseSrc = select.getFalseValue(); - if (auto load = trueSrc.getDefiningOp()) { - if (store.getDst() == load.getSrc() && load.getToken() == store.getDep()) - valueFromTMEM = kTrue; - } - if (auto load = falseSrc.getDefiningOp()) { - if (store.getDst() == load.getSrc() && load.getToken() == store.getDep()) - valueFromTMEM = valueFromTMEM == kTrue ? kUnknown : kFalse; - } - if (valueFromTMEM == kUnknown) { - return failure(); - } - Value pred = select.getCondition(); - // In case the false operand is overwriting, we need to negate the predicate - // (owerwrite when select would be false) - if (valueFromTMEM == kTrue) { - Value one = arith::ConstantIntOp::create(rewriter, select.getLoc(), 1, 1); - pred = arith::XOrIOp::create(rewriter, select.getLoc(), pred, one); - } - // Store the selected value with the updated predicate - Value overwritingValue = valueFromTMEM == kTrue ? falseSrc : trueSrc; - rewriter.replaceOpWithNewOp( - store, rewriter.getType(), store.getDst(), - store.getDep(), overwritingValue, pred); - return success(); - } -}; - -class RemoveUnusedTMEMLoad : public OpRewritePattern { -public: - using OpRewritePattern::OpRewritePattern; - - LogicalResult matchAndRewrite(ttng::TMEMLoadOp load, - PatternRewriter &rewriter) const override { - if (!load.getDep()) - return failure(); - if (!load.getResult().use_empty()) - return failure(); - rewriter.replaceAllUsesWith(load.getToken(), load.getDep()); - return success(); - } -}; - -// Load-store forwarding pattern. -class CombineTMEMLoadAndStore : public OpRewritePattern { -public: - using OpRewritePattern::OpRewritePattern; - - LogicalResult matchAndRewrite(ttng::TMEMStoreOp store, - PatternRewriter &rewriter) const override { - if (!store.getDep()) - return failure(); - auto load = store.getDep().getDefiningOp>(); - if (!load || load.getResult() != store.getSrc() || - load.getSrc() != store.getDst()) - return failure(); - rewriter.replaceOp(store, load.getToken()); - return success(); - } -}; - -class SinkTMEMLoad : public OpRewritePattern { -public: - using OpRewritePattern::OpRewritePattern; - - LogicalResult matchAndRewrite(ttng::TMEMLoadOp load, - PatternRewriter &rewriter) const override { - if (!load.getDep()) - return failure(); - auto forOp = load->getParentOfType(); - if (!forOp) { - return failure(); - } - DominanceInfo domInfo(forOp); - Operation *domOp = findNearestCommonDominator( - llvm::to_vector(load.getResult().getUsers()), domInfo); - if (!domOp || !domInfo.properlyDominates(load.getOperation(), domOp)) { - return failure(); - } - // Don't sink past potentially aliasing ops. - PostDominanceInfo postDomInfo(forOp); - SmallVector uses; - for (OpOperand &use : load.getToken().getUses()) - uses.push_back(&use); - if (!llvm::all_of(uses, [&](OpOperand *use) { - return postDomInfo.properlyPostDominates(use->getOwner(), domOp); - })) - return failure(); - // In order to not re-ordering multiple tmem load in a loop, don't sink if - // all the ops between the load and the domOp are tmem loads. - Operation *nextNode = load->getNextNode(); - while (auto tmemLoad = dyn_cast(nextNode)) { - nextNode = tmemLoad->getNextNode(); - } - if (domOp == nextNode) { - // The load wasn't moved. - return failure(); - } - rewriter.moveOpBefore(load, domOp); - Value newToken = sinkValueRedefinition(rewriter, load.getDep(), - load.getToken(), domOp->getBlock()); - if (newToken != load.getToken()) { - for (OpOperand *use : uses) - use->set(newToken); - } - return success(); - } -}; - -// Combine back TMEM alloc and store. This is equivalent but gives us a more -// canonical form to do further optimizations. -class CombineTMEMStoreAndAlloc : public OpRewritePattern { -public: - using OpRewritePattern::OpRewritePattern; - - LogicalResult matchAndRewrite(ttng::TMEMStoreOp store, - PatternRewriter &rewriter) const override { - if (!store.getDep()) - return failure(); - if (!matchPattern(store.getPred(), m_One())) - return failure(); - auto alloc = store.getDep().getDefiningOp(); - if (!alloc) - return failure(); - if (store.getDst() != alloc.getResult()) - return failure(); - if (alloc->getBlock() != store->getBlock()) - return failure(); - if (auto srcDef = store.getSrc().getDefiningOp()) { - if (alloc->getBlock() == srcDef->getBlock() && - alloc->isBeforeInBlock(srcDef)) - return failure(); - } - alloc.getSrcMutable().assign(store.getSrc()); - rewriter.replaceOp(store, alloc.getToken()); - return success(); - } -}; - -// Hoists a tmem alloc outside an if op like this: -// %0 = scf.if { -// %1, %token0 = tmem.alloc %init -// ... -// %2 = tmem.load %1, %token1 -// scf.yield %2 -// } else { -// scf.yield %init -// } -// -> -// %a, %token0 = tmem.alloc %init -// %token2 = scf.if { -// -// ... -// scf.yield %token1 -// } else { -// scf.yield %token0 -// } -// %2 = tmem.load %a, %token2 -class HoistTMEMAllocOutOfIf : public OpRewritePattern { -public: - using OpRewritePattern::OpRewritePattern; - - LogicalResult matchAndRewrite(ttng::TMEMAllocOp alloc, - PatternRewriter &rewriter) const override { - if (!alloc.getToken()) - return failure(); - Value init = alloc.getSrc(); - if (!init) - return failure(); - auto ifOp = dyn_cast(alloc->getParentOp()); - if (!ifOp || !ifOp.elseBlock()) - return failure(); - auto thenOp = ifOp.thenBlock()->getTerminator(); - auto elseOp = ifOp.elseBlock()->getTerminator(); - SmallVector yieldArgs; - for (auto [thenOperand, elseOperand] : - llvm::zip(thenOp->getOpOperands(), elseOp->getOpOperands())) { - auto load = thenOperand.get().getDefiningOp(); - if (!load || load.getSrc() != alloc.getResult()) - continue; - if (elseOperand.get() != init) - continue; - yieldArgs.push_back(thenOperand.getOperandNumber()); - } - if (yieldArgs.empty()) - return failure(); - // Since init is used in the else terminator we know that it dominates the - // if op. - alloc->moveBefore(ifOp); - rewriter.setInsertionPointAfter(ifOp); - for (int argNo : yieldArgs) { - auto load = - cast(thenOp->getOperand(argNo).getDefiningOp()); - auto newLoad = cast(rewriter.clone(*load)); - rewriter.modifyOpInPlace(ifOp, [&] { - ifOp->getResult(argNo).replaceAllUsesWith(newLoad.getResult()); - newLoad.getDepMutable().assign(ifOp->getResult(argNo)); - thenOp->setOperand(argNo, load.getToken()); - elseOp->setOperand(argNo, alloc.getToken()); - ifOp->getResult(argNo).setType(newLoad.getToken().getType()); - }); - } - return success(); - } -}; - -// Forward a TMEM load into the user allocation. -class TMEMLoadForwarding : public OpRewritePattern { -public: - using OpRewritePattern::OpRewritePattern; - - LogicalResult matchAndRewrite(ttng::TMEMAllocOp alloc, - PatternRewriter &rewriter) const override { - if (!alloc.getToken()) - return failure(); - Value init = alloc.getSrc(); - if (!init) - return failure(); - auto load = init.getDefiningOp(); - if (!load || !load->hasOneUse() || !load.getDep().hasOneUse()) - return failure(); - if (alloc.getType() != load.getSrc().getType()) - return failure(); - rewriter.replaceOp(alloc, {load.getSrc(), load.getDep()}); - return success(); - } -}; - -// Remove loop-carried tensor dependencies if they are fed immediately into a -// TMEM store by pulling the store into the previous iteration. -class RotateTMEMStoreInLoop : public OpRewritePattern { -public: - using OpRewritePattern::OpRewritePattern; - - LogicalResult matchAndRewrite(ttng::TMEMStoreOp store, - PatternRewriter &rewriter) const override { - if (!store.getDep()) - return failure(); - // Pattern match stores whose source comes from a loop region argument and - // whose predicate is loop-invariant. - scf::ForOp forOp = dyn_cast(store->getParentOp()); - if (!forOp || !forOp.isDefinedOutsideOfLoop(store.getPred()) || - !forOp.isDefinedOutsideOfLoop(store.getDst())) { - return failure(); - } - auto getAsLoopArg = [&](Value v) -> BlockArgument { - auto arg = dyn_cast(v); - if (arg && arg.getOwner() == forOp.getBody()) - return arg; - return {}; - }; - BlockArgument src = getAsLoopArg(store.getSrc()); - if (!src || !src.hasOneUse()) { - return failure(); - } - - // Check that rotating the store into the past won't violate any - // write-after-read dependencies. - BlockArgument storeTok = getAsLoopArg(store.getDep()); - if (!storeTok) - return failure(); - int tokArgNo = storeTok.getArgNumber() - 1; - - // Create two copies of the store: one before the loop, storing the initial - // value, and one before the yield, storing the value carried by the loop - // arg. - int argNo = src.getArgNumber() - 1; - Value initVal = forOp.getInitArgs()[argNo]; - rewriter.setInsertionPoint(forOp); - auto tokType = rewriter.getType(); - auto initStore = ttng::TMEMStoreOp::create( - rewriter, store.getLoc(), tokType, store.getDst(), - forOp.getInitArgs()[tokArgNo], initVal, store.getPred()); - forOp.getInitArgsMutable()[tokArgNo].assign(initStore.getToken()); - - auto yield = cast(forOp.getBody()->getTerminator()); - store.getToken().replaceAllUsesWith(forOp.getRegionIterArg(tokArgNo)); - rewriter.moveOpBefore(store, yield); - store.getDepMutable().assign(yield.getOperand(tokArgNo)); - yield.setOperand(tokArgNo, store.getToken()); - store.getSrcMutable().assign(yield.getOperand(argNo)); - - // Load from the tmem after the loop, and use it instead of the loop carried - // value. - rewriter.setInsertionPointAfter(forOp); - auto load = ttng::TMEMLoadOp::create( - rewriter, store.getLoc(), store.getSrc().getType(), tokType, - store.getDst(), forOp.getResult(tokArgNo)); - forOp->getResult(argNo).replaceAllUsesWith(load.getResult()); - // Loop carried value is no longer used, short-circuit it. - yield.setOperand(argNo, forOp.getRegionIterArg(argNo)); - return success(); - } -}; - -// Remove loop-carried tensor dependencies if they are the result of TMEM loads -// at the end of the loop by pushing the load into the next iteration. -class RotateTMEMLoadInLoop : public OpRewritePattern { -public: - using OpRewritePattern::OpRewritePattern; - - LogicalResult matchAndRewrite(ttng::TMEMLoadOp load, - PatternRewriter &rewriter) const override { - if (!load.getDep()) - return failure(); - // Pattern match loads whose results are only passed into the next iteration - // of a loop. - scf::ForOp forOp = dyn_cast(load->getParentOp()); - if (!forOp || !forOp.isDefinedOutsideOfLoop(load.getSrc()) || - !load.getResult().hasOneUse()) { - return failure(); - } - OpOperand &use = *load.getResult().use_begin(); - auto yield = dyn_cast(use.getOwner()); - if (!yield) - return failure(); - - // By rotating the load into the future, we are essentially merging the - // loop-carried tensor value into the same TMEM allocation as the load. - // Thus, they cannot be live at the same time. Check this by ensuring we - // won't clobber the memory. - - // 1. There are no aliasing stores between the load and the end of the loop. - if (!llvm::is_contained(load.getToken().getUsers(), yield)) - return failure(); - // 2. The TMEM variable is live into the loop with an undefined value. - int tokArgNo = load.getToken().use_begin()->getOperandNumber(); - Value initTok = forOp.getInitArgs()[tokArgNo]; - auto initAlloc = initTok.getDefiningOp(); - if (!initAlloc || initAlloc.getSrc()) - return failure(); - // TODO: 3. The live-in value of the TMEM variable is never read. - - // Create a store before the loop to write the initial value. - int argNo = use.getOperandNumber(); - Value initVal = forOp.getInitArgs()[argNo]; - rewriter.setInsertionPoint(forOp); - auto vTrue = arith::ConstantIntOp::create(rewriter, load.getLoc(), 1, 1); - auto tokType = rewriter.getType(); - auto initStore = ttng::TMEMStoreOp::create( - rewriter, load.getLoc(), tokType, load.getSrc(), initAlloc.getToken(), - initVal, vTrue); - forOp.getInitArgsMutable()[tokArgNo].assign(initStore.getToken()); - - // Move the load to the beginning of the loop to load the tensor value. - yield.setOperand(tokArgNo, load.getDep()); - rewriter.moveOpBefore(load, &forOp.getBody()->front()); - Value tokArg = forOp.getRegionIterArg(tokArgNo); - load.getDepMutable().assign(tokArg); - tokArg.replaceAllUsesExcept(load.getToken(), load); - forOp.getRegionIterArg(argNo).replaceAllUsesWith(load.getResult()); - - // Load from the tmem after the loop, and use it instead of the loop carried - // value. - rewriter.setInsertionPointAfter(forOp); - auto loadAfterLoop = ttng::TMEMLoadOp::create( - rewriter, load.getLoc(), load.getResult().getType(), tokType, - load.getSrc(), forOp.getResult(tokArgNo)); - forOp->getResult(argNo).replaceAllUsesWith(loadAfterLoop.getResult()); - // Loop carried value is no longer used, short-circuit it. - yield.setOperand(argNo, forOp.getRegionIterArg(argNo)); - return success(); - } -}; - -// Given an operation that uses a token, return its forwarded token. This -// assumes the memory variable is not loop carried. -static Value getTokenFromOp(Operation *op) { - if (auto mmaOp = dyn_cast>(op)) { - return mmaOp.getToken(); - } else if (auto loadOp = dyn_cast(op)) { - return loadOp.getToken(); - } else if (auto storeOp = dyn_cast(op)) { - return storeOp.getToken(); - } - assert(!isa(op) && "unexpected loop carried token"); - llvm_unreachable("unknown TMEM memory user"); -} - -// Find all the last uses of a memory variable in a loop body. This traces the -// token lattice to its leaves. -static void findLastMemoryUses(OpResult token, - SmallVectorImpl &lastUses, - DenseSet &seen) { - if (!seen.insert(token).second) - return; - if (token.use_empty()) { - lastUses.push_back(token); - return; - } - for (Operation *user : token.getUsers()) - findLastMemoryUses(cast(getTokenFromOp(user)), lastUses, seen); -} - -// Find the last uses of a memory variable, joining them into a single token if -// necessary. This token can be carried into the next loop iteration. -static Value joinLastMemoryUses(OpBuilder &b, Value token) { - SmallVector lastUses; - DenseSet seenTokens; - findLastMemoryUses(cast(token), lastUses, seenTokens); - assert(!lastUses.empty()); - - if (lastUses.size() == 1 && lastUses.front().getDefiningOp()->getBlock() == - token.getDefiningOp()->getBlock()) - return lastUses.front(); - // We can handle this case as needed. Right now it never happens. - llvm::report_fatal_error( - "FIXME: can't hoist TMEM alloc with multiple or conditional uses"); -} - -ttng::TMEMAllocOp hoistTMEMAlloc(TMEMTokenAllocOp alloc, scf::ForOp &forOp) { - OpBuilder builder(alloc); - builder.setInsertionPoint(forOp); - Value vTrue = arith::ConstantIntOp::create(builder, alloc.getLoc(), 1, 1); - auto src = alloc.getSrc(); - auto newAlloc = cast(builder.clone(*alloc)); - newAlloc.getSrcMutable().clear(); - - // By hoisting the allocation out of the loop, we need to turn the underlying - // memory variable into a loop-carried depdendency. - auto tokType = builder.getType(); - forOp = addIterArgsToLoop(builder, forOp, newAlloc.getToken()); - Value newTok = forOp.getRegionIterArgs().back(); - appendToForOpYield(forOp, joinLastMemoryUses(builder, alloc.getToken())); - - if (src != nullptr) { - builder.setInsertionPoint(alloc); - // Write the initial value of the allocation and replace the token. - auto initStoreOp = - ttng::TMEMStoreOp::create(builder, alloc.getLoc(), tokType, - newAlloc.getResult(), newTok, src, vTrue); - newTok = initStoreOp.getToken(); - } - alloc.replaceAllUsesWith(ValueRange{newAlloc.getResult(), newTok}); - alloc.erase(); - - return newAlloc; -} - -// Hoist invariant tmem_alloc. This could technically be done as general LICM -// but controlling tmem liveranga more precisley is likely to be important. -static void hoistInvariantInputs(Operation *mmaOp, scf::ForOp forOp) { - for (auto operand : mmaOp->getOperands()) { - if (forOp.isDefinedOutsideOfLoop(operand)) - continue; - auto tmemAllocOp = operand.getDefiningOp(); - if (!tmemAllocOp || tmemAllocOp.getType().getMutableMemory()) - continue; - assert(tmemAllocOp.getSrc()); - Value src = tmemAllocOp.getSrc(); - SmallVector opToHoist = {tmemAllocOp.getOperation()}; - // Also hoist simple unary elementwise that may have sinked into the loop. - while (Operation *defOp = src.getDefiningOp()) { - if (forOp.isDefinedOutsideOfLoop(src)) - break; - if (!(isPure(defOp) && defOp->getNumOperands() == 1)) - break; - opToHoist.push_back(defOp); - src = defOp->getOperand(0); - } - if (!forOp.isDefinedOutsideOfLoop(src)) - continue; - for (auto op : llvm::reverse(opToHoist)) { - forOp.moveOutOfLoop(op); - } - } -} -} // namespace - -struct HoistTMEMAlloc - : public impl::TritonGPUHoistTMEMAllocBase { - using impl::TritonGPUHoistTMEMAllocBase< - HoistTMEMAlloc>::TritonGPUHoistTMEMAllocBase; - - void runOnOperation() override { - ModuleOp m = getOperation(); - if (!hoistOutOfIf) { - SmallVector mmaOps; - m.walk([&](ttng::MMAv5OpInterface mmaOp) { mmaOps.push_back(mmaOp); }); - for (auto mmaOp : mmaOps) { - auto forOp = dyn_cast(mmaOp->getParentOp()); - if (!forOp) { - continue; - } - hoistInvariantInputs(mmaOp, forOp); - - // Only hoist the TMEM alloc feeding into the accumulator. Leave the - // ones for the scales in the loop. - auto alloc = mmaOp.getAccumulator().getDefiningOp(); - if (!alloc || alloc->getParentRegion() != mmaOp->getParentRegion()) { - continue; - } - hoistTMEMAlloc(alloc, forOp); - } - } - - mlir::RewritePatternSet patterns(&getContext()); - patterns.add(&getContext()); - if (hoistOutOfIf) { - patterns.add(&getContext()); - } - scf::ForOp::getCanonicalizationPatterns(patterns, &getContext()); - if (failed(applyPatternsGreedily(getOperation(), std::move(patterns)))) { - llvm_unreachable("Failed to hoist tmem_store"); - } - - // TODO: currently some code assumes that a mutable tmem alloc doesn't have - // an initial value. As a workaround we break up the op in order to keep - // this form for the downstream passes. We should remove this once the - // downstread passes are fixed. - m.walk([&](ttng::TMEMAllocOp alloc) { - if (alloc.getType().getMutableMemory() && alloc.getSrc()) { - OpBuilder builder(alloc); - builder.setInsertionPointAfter(alloc); - auto store = ttng::TMEMStoreOp::create( - builder, alloc.getLoc(), builder.getType(), - alloc.getResult(), alloc.getToken(), alloc.getSrc(), - arith::ConstantIntOp::create(builder, alloc.getLoc(), 1, 1)); - alloc.getToken().replaceAllUsesExcept(store.getToken(), store); - alloc.getSrcMutable().clear(); - } - }); - } -}; - -} // namespace gpu -} // namespace triton -} // namespace mlir diff --git a/third_party/mthreads/lib/Dialect/TritonGPU/Transforms/LayoutPropagationUtility.cpp b/third_party/mthreads/lib/Dialect/TritonGPU/Transforms/LayoutPropagationUtility.cpp deleted file mode 100644 index 8ab0a818dd..0000000000 --- a/third_party/mthreads/lib/Dialect/TritonGPU/Transforms/LayoutPropagationUtility.cpp +++ /dev/null @@ -1,49 +0,0 @@ -#include "triton/Dialect/TritonGPU/Transforms/LayoutPropagationUtility.h" -#include "triton/Dialect/Triton/IR/Dialect.h" -#include "triton/Dialect/TritonGPU/IR/Attributes.h" -#include "triton/Dialect/TritonGPU/IR/LinearLayoutConversions.h" -#include "triton/Dialect/TritonGPU/Transforms/Utility.h" -#include -#include - -namespace mlir::triton::gpu { - -std::optional> -inferSourceLoadLayout(const LinearLayout &dstLayout, Operation *defOp) { - if (!defOp) - return std::nullopt; - return inferSourceLoadLayout( - LinearEncodingAttr::get(defOp->getContext(), dstLayout), defOp); -} - -std::optional> -inferSourceLoadLayout(LinearEncodingAttr dstLayout, Operation *defOp) { - Attribute curLayout = dstLayout; - Operation *curOp = defOp; - while (curOp) { - if (isa(curOp)) - break; // Found the load op; we are done here. - - if (auto cvtOp = dyn_cast(curOp)) { - // For convert op we keep the current layout to push through further. - curOp = cvtOp.getSrc().getDefiningOp(); - } else { - if (curOp->getNumOperands() != 1) - break; - curLayout = inferSrcEncoding(curOp, curLayout); - curOp = curOp->getOperand(0).getDefiningOp(); - } - } - auto loadOp = dyn_cast_or_null(curOp); - if (!loadOp) - return std::nullopt; - auto loadType = dyn_cast(loadOp.getType()); - if (!loadType) - return std::nullopt; - - return std::make_pair( - loadOp, - toLinearLayout(loadType.getShape(), cast(curLayout))); -} - -} // namespace mlir::triton::gpu diff --git a/third_party/mthreads/lib/Dialect/TritonGPU/Transforms/Pipeliner/TestPipelineLowerLoop.cpp b/third_party/mthreads/lib/Dialect/TritonGPU/Transforms/Pipeliner/TestPipelineLowerLoop.cpp deleted file mode 100644 index 7602bb4765..0000000000 --- a/third_party/mthreads/lib/Dialect/TritonGPU/Transforms/Pipeliner/TestPipelineLowerLoop.cpp +++ /dev/null @@ -1,32 +0,0 @@ -#include "triton/Dialect/TritonGPU/IR/Dialect.h" -#include "triton/Dialect/TritonGPU/Transforms/Passes.h" -#include "triton/Dialect/TritonGPU/Transforms/PipeliningUtility.h" -#include "triton/Dialect/TritonGPU/Transforms/Schedule.h" -#include "triton/Dialect/TritonGPU/Transforms/Utility.h" - -using namespace mlir; -namespace tt = mlir::triton; -namespace ttg = mlir::triton::gpu; - -namespace mlir { -namespace triton { -namespace gpu { - -#define GEN_PASS_DEF_TRITONGPUTESTPIPELINELOWERLOOP -#include "triton/Dialect/TritonGPU/Transforms/Passes.h.inc" - -struct TestPipelineLowerLoop - : public impl::TritonGPUTestPipelineLowerLoopBase { - using impl::TritonGPUTestPipelineLowerLoopBase< - TestPipelineLowerLoop>::TritonGPUTestPipelineLowerLoopBase; - - void runOnOperation() override { - ModuleOp m = getOperation(); - - lowerLoops(m); - } -}; - -} // namespace gpu -} // namespace triton -} // namespace mlir diff --git a/third_party/mthreads/lib/Dialect/TritonGPU/Transforms/ReorderInstructions.cpp b/third_party/mthreads/lib/Dialect/TritonGPU/Transforms/ReorderInstructions.cpp deleted file mode 100644 index 456a40f48d..0000000000 --- a/third_party/mthreads/lib/Dialect/TritonGPU/Transforms/ReorderInstructions.cpp +++ /dev/null @@ -1,178 +0,0 @@ -#include "mlir/Analysis/SliceAnalysis.h" -#include "mlir/Dialect/SCF/IR/SCF.h" -#include "mlir/IR/BuiltinAttributes.h" -#include "mlir/IR/Dominance.h" -#include "mlir/IR/IRMapping.h" -#include "mlir/IR/Matchers.h" -#include "mlir/IR/PatternMatch.h" -#include "mlir/IR/Verifier.h" -#include "mlir/Interfaces/InferTypeOpInterface.h" -#include "mlir/Pass/Pass.h" -#include "mlir/Pass/PassManager.h" -#include "mlir/Support/LLVM.h" -#include "mlir/Support/LogicalResult.h" -#include "mlir/Transforms/GreedyPatternRewriteDriver.h" -#include "mlir/Transforms/Passes.h" -#include "mlir/Transforms/RegionUtils.h" -#include "triton/Analysis/Utility.h" -#include "triton/Dialect/Triton/IR/Dialect.h" -#include "triton/Dialect/TritonGPU/IR/Dialect.h" -#include "triton/Dialect/TritonGPU/Transforms/Passes.h" -#include "triton/Dialect/TritonGPU/Transforms/TritonGPUConversion.h" -#include "triton/Dialect/TritonGPU/Transforms/Utility.h" - -namespace mlir { -namespace triton { -namespace gpu { - -#define GEN_PASS_DEF_TRITONGPUREORDERINSTRUCTIONS -#include "triton/Dialect/TritonGPU/Transforms/Passes.h.inc" - -static bool willIncreaseRegisterPressure(Operation *op) { - if (isa(op)) - return true; - auto cvt = dyn_cast(op); - if (!cvt) - return false; - if (mlir::isa( - cvt.getType().getEncoding())) - return true; - return false; -} - -// Return true if it has side effects that are either unknown or writes. -static bool hasWriteSideEffect(Operation *op) { - auto effects = getEffectsRecursively(op); - if (!effects) - return false; - return llvm::any_of(*effects, [](MemoryEffects::EffectInstance effect) { - return !isa(effect.getEffect()); - }); -} - -// Return true if there is a write side effect on any path between start and end -// ops. This assumes start dominates end. -static bool crossWriteSideEffectingOp(Operation *start, Operation *end) { - auto ancestor = start->getBlock()->findAncestorOpInBlock(*end); - // Couldn't find an ancestor in the same block, conservatively assume true. - if (!ancestor) - return true; - Operation *nextOp = start->getNextNode(); - while (nextOp) { - if ((hasWriteSideEffect(nextOp))) - return true; - if (nextOp == ancestor) - return false; - nextOp = nextOp->getNextNode(); - } - assert(false && "op doesn't dominate other"); - return true; -} - -class TritonGPUReorderInstructionsPass - : public impl::TritonGPUReorderInstructionsBase< - TritonGPUReorderInstructionsPass> { -public: - TritonGPUReorderInstructionsPass() = default; - - Operation *getFirstUse(Operation *op) { - std::vector users; - for (auto user : op->getUsers()) { - if (Operation *ancestor = op->getBlock()->findAncestorOpInBlock(*user)) - users.push_back(ancestor); - } - auto minOpIt = - llvm::min_element(users, [](mlir::Operation *a, mlir::Operation *b) { - return a->isBeforeInBlock(b); - }); - return minOpIt != users.end() ? *minOpIt : nullptr; - } - - void runOnOperation() override { - ModuleOp m = getOperation(); - mlir::DominanceInfo dom(m); - // sink conversion after the last dealloc - // before the first use ancestor in its block - m.walk([&](triton::gpu::ConvertLayoutOp op) { - auto curr = mlir::Block::iterator(op); - auto end = op->getBlock()->end(); - for (; curr != end && &*curr != getFirstUse(op); curr++) - if (isa(&*curr)) - op->moveAfter(&*curr); - }); - // Sink conversions into loops when they will increase - // register pressure - DenseMap opToMove; - auto moveAfter = [](Operation *lhs, Operation *rhs) { - lhs->moveAfter(rhs); - }; - m.walk([&](Operation *op) { - if (!willIncreaseRegisterPressure(op)) - return; - auto user_begin = op->user_begin(); - auto user_end = op->user_end(); - if (std::distance(user_begin, user_end) != 1) - return; - if (user_begin->getParentOfType() == - op->getParentOfType()) - return; - opToMove.insert({op, *user_begin}); - }); - for (auto &kv : opToMove) - kv.first->moveBefore(kv.second); - // Move alloc(load) immediately after dependent load - m.walk([&](triton::gpu::LocalAllocOp op) { - if (!op.getSrc()) - return; - Operation *argOp = op.getSrc().getDefiningOp(); - if (!argOp) - return; - // Don't hoist alloc if the src is a scalar as this may increase smem - // pressure for no benefits. - if (isa(argOp)) - return; - moveAfter(op, argOp); - }); - // Move transpositions just after their definition - opToMove.clear(); - m.walk([&](triton::TransposeOpInterface op) { - Operation *argOp = op.getSrc().getDefiningOp(); - if (!argOp) - return; - moveAfter(op, argOp); - }); - // Move `dot` operand so that conversions to opIdx=1 happens after - // conversions to opIdx=0 - m.walk([&](triton::gpu::LocalLoadOp op) { - auto dstEncoding = mlir::dyn_cast( - op.getType().getEncoding()); - if (!dstEncoding) - return; - int opIdx = dstEncoding.getOpIdx(); - if (opIdx != 1) - return; - if (!op->hasOneUse()) - return; - auto dotUser = dyn_cast(*op->user_begin()); - if (!dotUser) - return; - auto AOp = - dotUser.getOperand(0).getDefiningOp(); - if (!AOp) - return; - // Check that the conversion to OpIdx=1 happens before and can be moved - // after the conversion to OpIdx=0. - if (!dom.dominates(op.getOperation(), AOp.getOperation())) - return; - if (crossWriteSideEffectingOp(op, AOp)) - return; - moveAfter(op, AOp); - }); - return; - } -}; - -} // namespace gpu -} // namespace triton -} // namespace mlir diff --git a/third_party/mthreads/lib/Dialect/TritonGPU/Transforms/WarpSpecialization/PartitionBuilder.cpp b/third_party/mthreads/lib/Dialect/TritonGPU/Transforms/WarpSpecialization/PartitionBuilder.cpp deleted file mode 100644 index 8d18c1fab1..0000000000 --- a/third_party/mthreads/lib/Dialect/TritonGPU/Transforms/WarpSpecialization/PartitionBuilder.cpp +++ /dev/null @@ -1,36 +0,0 @@ -#include "triton/Dialect/TritonGPU/Transforms/PartitionBuilder.h" -#include "triton/Dialect/TritonGPU/Transforms/Partition.h" -#include "triton/Dialect/TritonGPU/Transforms/PipeliningUtility.h" - -using namespace mlir; -using namespace triton; -using namespace triton::gpu; - -Value PartitionBuilder::intCst(int value, unsigned width) { - return create(value, width); -} - -Value PartitionBuilder::boolCst(bool value) { - return intCst(value, /*width=*/1); -} - -void PartitionBuilder::assignPartition(Operation *op, Partition &partition) { - setPartition(op, &partition); -} - -StageCluster triton::gpu::getStageCluster(Operation *op) { - auto stageAttr = op->getAttrOfType(kLoopStageAttrName); - auto clusterAttr = op->getAttrOfType(kLoopClusterAttrName); - if (!stageAttr || !clusterAttr) - return std::nullopt; - return std::make_pair(stageAttr.getInt(), clusterAttr.getInt()); -} - -void triton::gpu::setStageCluster(OpBuilder &b, Operation *op, - StageCluster stageCluster) { - if (stageCluster) { - op->setAttr(kLoopStageAttrName, b.getI32IntegerAttr(stageCluster->first)); - op->setAttr(kLoopClusterAttrName, - b.getI32IntegerAttr(stageCluster->second)); - } -} diff --git a/third_party/mthreads/lib/Dialect/TritonInstrument/CMakeLists.txt b/third_party/mthreads/lib/Dialect/TritonInstrument/CMakeLists.txt deleted file mode 100644 index 9f57627c32..0000000000 --- a/third_party/mthreads/lib/Dialect/TritonInstrument/CMakeLists.txt +++ /dev/null @@ -1,2 +0,0 @@ -add_subdirectory(IR) -add_subdirectory(Transforms) diff --git a/third_party/mthreads/lib/Dialect/TritonInstrument/IR/CMakeLists.txt b/third_party/mthreads/lib/Dialect/TritonInstrument/IR/CMakeLists.txt deleted file mode 100644 index 6b39e076d6..0000000000 --- a/third_party/mthreads/lib/Dialect/TritonInstrument/IR/CMakeLists.txt +++ /dev/null @@ -1,14 +0,0 @@ -add_triton_library(TritonInstrumentIR - Dialect.cpp - FunctionBuilder.cpp - Ops.cpp - Utility.cpp - - DEPENDS - TritonInstrumentTableGen - - LINK_LIBS PUBLIC - MLIRIR - TritonIR - TritonGPUIR -) diff --git a/third_party/mthreads/lib/Dialect/TritonInstrument/IR/Dialect.cpp b/third_party/mthreads/lib/Dialect/TritonInstrument/IR/Dialect.cpp deleted file mode 100644 index d00906f30f..0000000000 --- a/third_party/mthreads/lib/Dialect/TritonInstrument/IR/Dialect.cpp +++ /dev/null @@ -1,17 +0,0 @@ -#include "triton/Dialect/Triton/IR/Dialect.h" -#include "mlir/IR/DialectImplementation.h" -#include "mlir/IR/OpImplementation.h" -#include "triton/Dialect/Triton/IR/Interfaces.h" -#include "triton/Dialect/TritonGPU/IR/Dialect.h" -#include "triton/Dialect/TritonInstrument/IR/Dialect.h" - -#include "triton/Dialect/TritonInstrument/IR/Dialect.cpp.inc" -using namespace mlir::triton::instrument; - -void TritonInstrumentDialect::initialize() { - addOperations< -#define GET_OP_LIST -#include "triton/Dialect/TritonInstrument/IR/Ops.cpp.inc" - >(); - addInterfaces(); -} diff --git a/third_party/mthreads/lib/Dialect/TritonInstrument/IR/Ops.cpp b/third_party/mthreads/lib/Dialect/TritonInstrument/IR/Ops.cpp deleted file mode 100644 index 823cc8649b..0000000000 --- a/third_party/mthreads/lib/Dialect/TritonInstrument/IR/Ops.cpp +++ /dev/null @@ -1,8 +0,0 @@ -#include "triton/Dialect/TritonGPU/IR/Dialect.h" -#include "triton/Dialect/TritonInstrument/IR/Dialect.h" -#include "triton/Dialect/TritonInstrument/IR/Utility.h" - -#define GET_OP_CLASSES -#include "triton/Dialect/TritonInstrument/IR/Ops.cpp.inc" - -#include "triton/Dialect/TritonInstrument/IR/OpsEnums.cpp.inc" diff --git a/third_party/mthreads/lib/Dialect/TritonInstrument/Transforms/CMakeLists.txt b/third_party/mthreads/lib/Dialect/TritonInstrument/Transforms/CMakeLists.txt deleted file mode 100644 index 62116e5927..0000000000 --- a/third_party/mthreads/lib/Dialect/TritonInstrument/Transforms/CMakeLists.txt +++ /dev/null @@ -1,16 +0,0 @@ -add_triton_library(TritonInstrumentTransforms - ConcurrencySanitizer.cpp - - DEPENDS - TritonInstrumentTransformsIncGen - - LINK_LIBS PUBLIC - MLIRTransforms - MLIRTransformUtils - TritonIR - TritonGPUIR - TritonNvidiaGPUIR - TritonToTritonGPU - TritonInstrumentIR - MLIRTransformUtils -) diff --git a/third_party/mthreads/lib/Dialect/TritonNvidiaGPU/CMakeLists.txt b/third_party/mthreads/lib/Dialect/TritonNvidiaGPU/CMakeLists.txt deleted file mode 100644 index 9f57627c32..0000000000 --- a/third_party/mthreads/lib/Dialect/TritonNvidiaGPU/CMakeLists.txt +++ /dev/null @@ -1,2 +0,0 @@ -add_subdirectory(IR) -add_subdirectory(Transforms) diff --git a/third_party/mthreads/lib/Dialect/TritonNvidiaGPU/IR/CMakeLists.txt b/third_party/mthreads/lib/Dialect/TritonNvidiaGPU/IR/CMakeLists.txt deleted file mode 100644 index c7a6bfa557..0000000000 --- a/third_party/mthreads/lib/Dialect/TritonNvidiaGPU/IR/CMakeLists.txt +++ /dev/null @@ -1,15 +0,0 @@ -add_triton_library(TritonNvidiaGPUIR - Dialect.cpp - TensorMemoryUtils.cpp - Ops.cpp - - DEPENDS - TritonNvidiaGPUTableGen - TritonNvidiaGPUAttrDefsIncGen - TritonNvidiaGPUOpInterfacesIncGen - TritonNvidiaGPUTypesIncGen - - LINK_LIBS PUBLIC - TritonIR - TritonGPUIR -) diff --git a/third_party/mthreads/lib/Dialect/TritonNvidiaGPU/Transforms/CMakeLists.txt b/third_party/mthreads/lib/Dialect/TritonNvidiaGPU/Transforms/CMakeLists.txt deleted file mode 100644 index 7715cc9861..0000000000 --- a/third_party/mthreads/lib/Dialect/TritonNvidiaGPU/Transforms/CMakeLists.txt +++ /dev/null @@ -1,25 +0,0 @@ -add_triton_library(TritonNvidiaGPUTransforms - CheckMatmulTwoCTAs.cpp - FenceInsertion.cpp - InterleaveTMem.cpp - MMALowering.cpp - OptimizeDescriptorEncoding.cpp - OptimizeTMemLayouts.cpp - PlanCTA.cpp - PromoteLHSToTMem.cpp - ProxyFenceInsertion.cpp - RemoveTMEMTokens.cpp - TensorMemoryAllocation.cpp - TMALowering.cpp - TMAUtilities.cpp - - DEPENDS - TritonNvidiaGPUTransformsIncGen - - LINK_LIBS PUBLIC - TritonIR - TritonGPUIR - TritonGPUTransforms - TritonNvidiaGPUIR - MLIRTransformUtils -) diff --git a/third_party/mthreads/lib/Dialect/TritonNvidiaGPU/Transforms/CheckMatmulTwoCTAs.cpp b/third_party/mthreads/lib/Dialect/TritonNvidiaGPU/Transforms/CheckMatmulTwoCTAs.cpp deleted file mode 100644 index c5b1ddf37a..0000000000 --- a/third_party/mthreads/lib/Dialect/TritonNvidiaGPU/Transforms/CheckMatmulTwoCTAs.cpp +++ /dev/null @@ -1,63 +0,0 @@ -#include "triton/Dialect/TritonNvidiaGPU/IR/Dialect.h" -#include "triton/Dialect/TritonNvidiaGPU/Transforms/Passes.h" - -#include "mlir/IR/BuiltinAttributes.h" -#include "mlir/IR/BuiltinOps.h" -#include "mlir/IR/Diagnostics.h" -#include "mlir/IR/Visitors.h" - -namespace ttng = mlir::triton::nvidia_gpu; - -namespace mlir::triton::nvidia_gpu { - -#define GEN_PASS_DEF_TRITONNVIDIAGPUCHECKMATMULTWOCTAPASS -#include "triton/Dialect/TritonNvidiaGPU/Transforms/Passes.h.inc" - -namespace { - -class TritonNvidiaGPUCheckMatmulTwoCTAPass - : public impl::TritonNvidiaGPUCheckMatmulTwoCTAPassBase< - TritonNvidiaGPUCheckMatmulTwoCTAPass> { -public: - using impl::TritonNvidiaGPUCheckMatmulTwoCTAPassBase< - TritonNvidiaGPUCheckMatmulTwoCTAPass>:: - TritonNvidiaGPUCheckMatmulTwoCTAPassBase; - - void runOnOperation() override { - ModuleOp mod = getOperation(); - Operation *firstMatmul = nullptr; - bool firstTwoCTA = false; - - WalkResult result = mod.walk([&](ttng::TCGen5MMAOp op) { - bool currentTwoCTA = op.getTwoCtas(); - if (!firstMatmul) { - firstMatmul = op; - firstTwoCTA = currentTwoCTA; - return WalkResult::advance(); - } - if (currentTwoCTA != firstTwoCTA) { - auto diag = op.emitError() - << "inconsistent two_ctas setting across matmuls; " - "expected all matmuls to " - << (firstTwoCTA ? "enable" : "disable") << " two_ctas."; - diag.attachNote(firstMatmul->getLoc()) - << "first matmul here has two_ctas=" - << (firstTwoCTA ? "true" : "false") << "."; - return WalkResult::interrupt(); - } - return WalkResult::advance(); - }); - - if (result.wasInterrupted()) { - signalPassFailure(); - return; - } - - bool twoCTAValue = firstMatmul ? firstTwoCTA : false; - mod->setAttr(AttrTwoCTAsName, BoolAttr::get(mod.getContext(), twoCTAValue)); - } -}; - -} // namespace - -} // namespace mlir::triton::nvidia_gpu diff --git a/third_party/mthreads/lib/Target/CMakeLists.txt b/third_party/mthreads/lib/Target/CMakeLists.txt deleted file mode 100644 index 39d31dc9b5..0000000000 --- a/third_party/mthreads/lib/Target/CMakeLists.txt +++ /dev/null @@ -1 +0,0 @@ -add_subdirectory(LLVMIR) diff --git a/third_party/mthreads/lib/Target/LLVMIR/CMakeLists.txt b/third_party/mthreads/lib/Target/LLVMIR/CMakeLists.txt deleted file mode 100644 index 88a265cd01..0000000000 --- a/third_party/mthreads/lib/Target/LLVMIR/CMakeLists.txt +++ /dev/null @@ -1,31 +0,0 @@ -add_triton_library(TritonLLVMIR - LLVMDIScope.cpp - LLVMDILocalVariable.cpp - LLVMIRBreakPhiStruct.cpp - LLVMDIUtils.cpp - - DEPENDS - LLVMIRIncGen - - LINK_LIBS - ${CMAKE_DL_LIBS} - PUBLIC - MLIRArithToLLVM - MLIRBuiltinToLLVMIRTranslation - MLIRIndexToLLVM - MLIRIR - MLIRLLVMDialect - MLIRNVVMToLLVM - MLIRLLVMToLLVMIRTranslation - MLIRNVVMToLLVMIRTranslation - MLIRROCDLToLLVMIRTranslation - MLIRSCFToControlFlow - MLIRSupport - MLIRTargetLLVMIRExport - TritonGPUToLLVM - ) - -set_source_files_properties( - LLVMIRTranslation.cpp - PROPERTIES - COMPILE_FLAGS "-D__BUILD_DIR__=\\\"${CMAKE_BINARY_DIR}\\\"") diff --git a/third_party/mthreads/lib/Target/LLVMIR/LLVMIRBreakPhiStruct.cpp b/third_party/mthreads/lib/Target/LLVMIR/LLVMIRBreakPhiStruct.cpp deleted file mode 100644 index a3c6d69959..0000000000 --- a/third_party/mthreads/lib/Target/LLVMIR/LLVMIRBreakPhiStruct.cpp +++ /dev/null @@ -1,60 +0,0 @@ -//===----------------------------------------------------------------------===// -/// Implements a trivial pass breaking up 1 level deep structure in phi nodes. -/// This handles the common case generated by Triton and allow better -/// optimizations down the compiler pipeline. -//===----------------------------------------------------------------------===// -#include "LLVMPasses.h" -#include "llvm/IR/IRBuilder.h" -#include "llvm/IR/Instructions.h" - -using namespace llvm; - -static bool processPhiStruct(PHINode *phiNode) { - StructType *STy = dyn_cast(phiNode->getType()); - if (!STy) - return false; - IRBuilder<> builder(phiNode); - unsigned numOperands = phiNode->getNumIncomingValues(); - unsigned numScalarEl = STy->getNumElements(); - Value *newStruct = UndefValue::get(STy); - builder.SetInsertPoint(phiNode->getParent()->getFirstNonPHIIt()); - llvm::IRBuilderBase::InsertPoint insertInsertPt = builder.saveIP(); - for (unsigned i = 0; i < numScalarEl; i++) { - builder.SetInsertPoint(phiNode); - PHINode *newPhiNode = - builder.CreatePHI(STy->getElementType(i), numOperands); - for (unsigned j = 0; j < numOperands; ++j) { - Value *operand = phiNode->getIncomingValue(j); - builder.SetInsertPoint(phiNode->getIncomingBlock(j)->getTerminator()); - newPhiNode->addIncoming(builder.CreateExtractValue(operand, i), - phiNode->getIncomingBlock(j)); - } - builder.restoreIP(insertInsertPt); - newStruct = builder.CreateInsertValue(newStruct, newPhiNode, i); - insertInsertPt = builder.saveIP(); - } - phiNode->replaceAllUsesWith(newStruct); - return true; -} - -static bool runOnFunction(Function &F) { - bool Changed = false; - SmallVector PhiNodes; - for (BasicBlock &BB : F) { - for (Instruction &inst : BB) { - if (PHINode *phiNode = dyn_cast(&inst)) { - Changed |= processPhiStruct(phiNode); - continue; - } - break; - } - } - return Changed; -} - -PreservedAnalyses BreakStructPhiNodesPass::run(Function &F, - FunctionAnalysisManager &AM) { - - bool b = runOnFunction(F); - return b ? PreservedAnalyses::none() : PreservedAnalyses::all(); -} diff --git a/third_party/mthreads/lib/Target/LLVMIR/LLVMPasses.h b/third_party/mthreads/lib/Target/LLVMIR/LLVMPasses.h deleted file mode 100644 index 1dcdb2992c..0000000000 --- a/third_party/mthreads/lib/Target/LLVMIR/LLVMPasses.h +++ /dev/null @@ -1,16 +0,0 @@ -#include "llvm/IR/PassManager.h" -#include "llvm/Pass.h" -#include "llvm/Support/CodeGen.h" - -namespace llvm { - -// Pass to pre-process LLVM IR before optimization and break up phi of struct. -// Breaking up those phis into elementary types allows better optimizations -// downstream. -struct BreakStructPhiNodesPass : PassInfoMixin { - PreservedAnalyses run(Function &F, FunctionAnalysisManager &AM); - - static StringRef name() { return "BreakStructPhiNodesPass"; } -}; - -} // namespace llvm diff --git a/third_party/mthreads/lib/Tools/CMakeLists.txt b/third_party/mthreads/lib/Tools/CMakeLists.txt deleted file mode 100644 index 611468b9a2..0000000000 --- a/third_party/mthreads/lib/Tools/CMakeLists.txt +++ /dev/null @@ -1,13 +0,0 @@ -add_triton_library(TritonTools - GenericSwizzling.cpp - LayoutUtils.cpp - LinearLayout.cpp - PluginUtils.cpp - - DEPENDS - - LINK_LIBS PUBLIC - MLIRIR - MLIRLLVMDialect - f2reduce -) diff --git a/third_party/mthreads/lib/Tools/GenericSwizzling.cpp b/third_party/mthreads/lib/Tools/GenericSwizzling.cpp deleted file mode 100644 index fedd25a3c3..0000000000 --- a/third_party/mthreads/lib/Tools/GenericSwizzling.cpp +++ /dev/null @@ -1,713 +0,0 @@ -#include "triton/Tools/GenericSwizzling.h" - -#include "third_party/f2reduce/f2reduce.h" -#include "triton/Tools/LayoutUtils.h" -#include "triton/Tools/LinearLayout.h" -#include "llvm/ADT/DenseSet.h" -#include "llvm/ADT/STLExtras.h" -#include "llvm/ADT/SmallVector.h" -#include "llvm/Support/Debug.h" -#include "llvm/Support/MathExtras.h" - -#define DEBUG_TYPE "generic-swizzling" -#define DBGS() (llvm::dbgs() << "[" DEBUG_TYPE "]: ") - -#if defined(_MSC_VER) && !defined(__clang__) -// from https://gist.github.com/pps83/3210a2f980fd02bb2ba2e5a1fc4a2ef0 -#include - -static int __builtin_ctzll(unsigned long long x) { - unsigned long r; - _BitScanForward64(&r, x); - return static_cast(r); -} - -#endif - -void printBasis(const llvm::SmallVector &basis, - const std::string &name) { - llvm::errs() << name << ": "; - for (int32_t b : basis) - llvm::errs() << b << " "; - llvm::errs() << "\n"; -} - -using namespace mlir; -using namespace mlir::triton; - -namespace { - -// Goes from bases of the form [[1], [2], [4], [8]] to [1, 2, 4, 8] -SmallVector flatten(const LinearLayout &ll, StringAttr dim) { - assert(ll.getNumOutDims() == 1); - auto outDim = *ll.getOutDimNames().begin(); - SmallVector vec; - for (int i = 0; i < ll.getInDimSizeLog2(dim); ++i) - vec.push_back(ll.getBasis(dim, i, outDim)); - return vec; -}; - -SmallVector removeZeros(ArrayRef vec) { - SmallVector result; - for (int32_t r : vec) { - if (r != 0) { - result.push_back(r); - } - } - return result; -} - -// [1, 2, 4, 8] -> [[1], [2], [4], [8]] -std::vector> unflatten(ArrayRef basis) { - std::vector> unflattened; - for (int32_t b : basis) - unflattened.push_back({b}); - return unflattened; -} - -// Compute the nullspace basis of `vectors` -SmallVector nullspaceBasis(ArrayRef vectors, int32_t dim) { - // Solve A^T x = 0, where A is the matrix of vectors - // To do this, we form a matrix where each vector is a row - const int32_t nRows = vectors.size(); - auto mat = std::make_unique(nRows); - for (int i = 0; i < nRows; ++i) - mat[i] = static_cast(vectors[i]); - f2reduce::inplace_rref_strided(mat.get(), /*rows=*/nRows, /*cols=*/dim, - /*stride=*/1); - - llvm::SmallDenseSet pivotCols; - for (int32_t r = 0; r < nRows; ++r) - if (mat[r]) - pivotCols.insert(__builtin_ctzll(mat[r])); - - SmallVector basis; - for (int32_t freeCol = 0; freeCol < dim; ++freeCol) { - if (!pivotCols.contains(freeCol)) { - uint64_t vec = 1ull << freeCol; - for (int32_t r = 0; r < nRows; ++r) - if (mat[r] & (1ull << freeCol)) { - const int32_t pivot = __builtin_ctzll(mat[r]); - vec ^= (1ull << pivot); - } - basis.push_back(static_cast(vec)); - } - } - return basis; -} - -// Find the smallest tile that we can read and write to smem -// without sacrificing vectorisation and split it into its own -// `reps` dimension -LinearLayout buildReps(MLIRContext *ctx, const LinearLayout &src, - const LinearLayout &dst, const LinearLayout &smem, - int32_t leaveReps) { - auto kVec = StringAttr::get(ctx, "vector"); - auto kBank = StringAttr::get(ctx, "bank"); - auto kSegment = StringAttr::get(ctx, "segment"); - auto kReps = StringAttr::get(ctx, "reps"); - auto kReg = StringAttr::get(ctx, "register"); - // A basis is a rep if: - // 1) It is in registers in both src and dst - // 2) It is in the segment of smem (i.e., is not part of just one - // load/store) - SetVector srcRegs(llvm::from_range_t{}, flatten(src, kReg)); - SetVector dstRegs(llvm::from_range_t{}, flatten(dst, kReg)); - SetVector smemSegment(llvm::from_range_t{}, flatten(smem, kSegment)); - SetVector segment; - SetVector reps; - for (auto s : smemSegment) { - // Do not move the first leaveReps bases from reps to segment - // as we need them to vectorise the instructions (think .x2 and .x4 in - // ldmatrix) - if (srcRegs.contains(s) && dstRegs.contains(s)) { - if (leaveReps > 0) { - leaveReps--; - segment.insert(s); - } else { - reps.insert(s); - } - } else { - segment.insert(s); - } - } - - auto smemReps = LinearLayout({{kVec, smem.getBases().lookup(kVec)}, - {kBank, smem.getBases().lookup(kBank)}, - {kSegment, unflatten(to_vector(segment))}, - {kReps, unflatten(to_vector(reps))}}, - smem.getOutDims(), - /*requireSurjective=*/true); - return smemReps; -} - -SmallVector computeSegment(const SmallVector &bankSrc, - const SmallVector &bankDst, - int32_t dim, int32_t lenSegment) { - llvm::SmallDenseSet setSrc(bankSrc.begin(), bankSrc.end()); - llvm::SmallDenseSet setDst(bankDst.begin(), bankDst.end()); - // Remove the 0 as it's not a basis - setSrc.erase(0); - setDst.erase(0); - - SmallVector segment; - for (int32_t b = 0; b < dim; ++b) - if (!setSrc.contains(1 << b) && !setDst.contains(1 << b)) - segment.push_back(1 << b); - if (segment.size() >= lenSegment) { - segment.resize(lenSegment); - return segment; - } - - // A and B are the difference sets - SmallVector A, B; - for (int32_t v : setSrc) - if (!setDst.contains(v)) - A.push_back(v); - for (int32_t v : setDst) - if (!setSrc.contains(v)) - B.push_back(v); - if (A.size() > B.size()) { - std::swap(A, B); - } - llvm::sort(A); - llvm::sort(B); - // A is the smaller set now - auto logBankConflicts = std::min( - std::max(0, lenSegment - A.size() - segment.size()), A.size()); - // Conflict-free - for (int i = logBankConflicts; i < A.size(); ++i) - segment.push_back(A[i] ^ B[i]); - // Write conflicts - segment.append(A.begin(), A.begin() + logBankConflicts); - // Read conflicts - segment.append(B.begin(), B.begin() + logBankConflicts); - - if (segment.size() > lenSegment) - segment.resize(lenSegment); - return segment; -} - -SmallVector complementBasis(ArrayRef basis, int32_t dim) { - const int32_t nRows = basis.size(); - auto mat = std::make_unique(nRows); - for (int r = 0; r < nRows; ++r) - mat[r] = static_cast(basis[r]); - - f2reduce::inplace_rref_strided(mat.get(), /*rows=*/nRows, - /*cols=*/dim, /*stride=*/1); - - llvm::SmallDenseSet pivotCols; - for (int r = 0; r < nRows; ++r) { - if (mat[r]) { - pivotCols.insert(__builtin_ctzll(mat[r])); // leading-1 position - } - } - - SmallVector comp; - for (int i = 0; i < dim; ++i) - if (!pivotCols.contains(i)) - comp.push_back(1 << i); - - return comp; -} -} // namespace - -namespace mlir::triton::gpu { - -SmallVector intersectionBasis(ArrayRef b1, - ArrayRef b2, int32_t dim) { - // If needed to be generic, this can be done computing - // nullspaceBasis(concat(nullspaceBasis(b1), nullspaceBasis(b2))) - // but doing this returns the bases in an arbitrary order! - auto isPowerOf2 = [](int32_t x) { return llvm::isPowerOf2_32(x); }; - bool powerOf2 = llvm::all_of(b1, isPowerOf2) && llvm::all_of(b2, isPowerOf2); - if (powerOf2) { - SmallVector result; - // Heuristic: We choose to retain the order relative to b1 - SetVector set2(b2.begin(), b2.end()); - for (int32_t b : b1) { - if (b != 0 && set2.contains(b)) { - result.push_back(b); - } - } - return result; - } else { - auto ns1 = nullspaceBasis(b1, dim); - auto ns2 = nullspaceBasis(b2, dim); - auto joint = llvm::to_vector(llvm::concat(ns1, ns2)); - return nullspaceBasis(joint, dim); - } -} - -std::pair bankConflicts(ArrayRef tileSrc, - ArrayRef tileDst, - const LinearLayout &smem) { - auto *ctx = smem.getOutDimNames().begin()->getContext(); - auto smemFlat = smem.flattenOuts(); - auto inDim = *smem.getInDimNames().begin(); - // Look at the intersection between the segment bases and the tile bases - // We don't need to intersect with the bases that covert the bank (as in - // the first 32 / bitwidth bases) because if we hit any of those broadcasting - // will avoid the bank conflict - auto segment = StringAttr::get(ctx, "segment"); - auto segmentBases = flatten(smemFlat, segment); - - int32_t rank = smem.getTotalOutDimSizeLog2(); - // compute conflicts - int write = 1 << intersectionBasis(segmentBases, tileSrc, rank).size(); - int read = 1 << intersectionBasis(segmentBases, tileDst, rank).size(); - return {read - 1, write - 1}; -} - -std::pair bankConflictsLdSt(const LinearLayout &src, - const LinearLayout &dst, - const LinearLayout &smem, - int32_t bitwidth) { - auto srcFlat = src.flattenOuts(); - auto dstFlat = dst.flattenOuts(); - auto *ctx = smem.getOutDimNames().begin()->getContext(); - auto S = [ctx](StringRef str) { return StringAttr::get(ctx, str); }; - auto kVec = S("vector"); - auto srcLane = flatten(srcFlat, S("lane")); - auto dstLane = flatten(dstFlat, S("lane")); - auto log2Vec = - llvm::Log2_32(std::max(smem.getInDimSize(kVec) * bitwidth / 32, 1)); - srcLane.resize(srcLane.size() - log2Vec); - dstLane.resize(dstLane.size() - log2Vec); - return bankConflicts(srcLane, dstLane, smem); -} - -int bankConflictsMemDesc(const LinearLayout ®, const LinearLayout &smem, - int32_t bitwidth) { - auto *ctx = smem.getInDimNames().begin()->getContext(); - auto S = [ctx](StringRef str) { return StringAttr::get(ctx, str); }; - - assert(smem.hasInDim(S("offset")) && "shared layout must have an offset dim"); - assert(reg.hasInDim(S("register")) && - "register layout must have a register dim"); - auto regNoBroadcast = actionRemoveBroadcastedRegs(reg).apply(reg); - auto regToShared = regNoBroadcast.invertAndCompose(smem); - auto [elemsPerVec, permutation] = - largestVectorisation(ctx, regToShared, bitwidth); - regNoBroadcast = permutation.apply(regNoBroadcast); - - int32_t vecSize = elemsPerVec; - int32_t bankSize = - std::min(32 * 32 / (vecSize * bitwidth), smem.getTotalInDimSize()); - int32_t segmentSize = smem.getTotalInDimSize() / (bankSize * vecSize); - SmallVector> newInDims = { - {S("vector"), vecSize}, - {S("bank"), bankSize}, - {S("segment"), segmentSize}, - }; - auto smemReshaped = smem.reshapeIns(newInDims); - return bankConflictsLdSt(regNoBroadcast, regNoBroadcast, smemReshaped, - bitwidth) - .first; -} - -std::optional> optimalSwizzlingTile( - const LinearLayout &a, const LinearLayout &b, int32_t nRegA, int32_t nRegB, - ArrayRef laneIdTileA, ArrayRef laneIdTileB) { - // For now se just implement the .v4 variants for all the instructions - // We could generalise this in the future - assert(nRegA + laneIdTileA.size() == nRegB + laneIdTileB.size()); - // normalise nRegA >= nRegB - if (nRegA < nRegB) { - return optimalSwizzlingTile(b, a, nRegB, nRegA, laneIdTileB, laneIdTileA); - } - assert(nRegA >= nRegB); - - auto *ctx = a.getInDimNames().begin()->getContext(); - auto kReg = StringAttr::get(ctx, "register"); - auto kLane = StringAttr::get(ctx, "lane"); - auto dim = a.getTotalOutDimSizeLog2(); - // map from b to a - LinearLayout cvt = b.invertAndCompose(a); - - // The contiguous tile of ld.shared.b32.v4 for a packed element of size - // bitwidth is composed of 128/bitwidth register elements - // The contiguous tile of ldmatrix.v4 for a packed element of size bitwidth - // is composed of 32/bitwidth register elements and the bases 0, 1st as given - // by the laneAddr - // The contiguous tile of ldmatrix.v4.trans for a packed element of size 16 - // is composed of the bases 2, 3, 4th as given by the laneAddr - - // Note that for register elements, we can choose any register basis we want, - // but the lane bases are fixed - - // In this function, we compute a tile (set of bases) such that it matches - // the tiles of A and B - - auto regA = flatten(a, kReg); - auto regB = flatten(b, kReg); - auto laneA = flatten(a, kLane); - auto laneB = flatten(b, kLane); - - // Compute the number of registers that start the tile - SmallVector vbasis = intersectionBasis(regA, regB, dim); - // We need to have at least nRegB vectorisation - if (vbasis.size() < nRegB) { - return std::nullopt; - } - vbasis.resize(nRegB); - - auto index = [](ArrayRef lane, ArrayRef laneIdTile) { - SmallVector ret; - for (auto id : laneIdTile) { - ret.push_back(lane[id]); - } - return ret; - }; - auto laneTileA = index(laneA, laneIdTileA); - auto laneTileB = index(laneB, laneIdTileB); - - // We need the tiles to be contiguous - auto isZero = [](int32_t b) { return b == 0; }; - if (llvm::any_of(laneTileA, isZero) || llvm::any_of(laneTileB, isZero)) { - return std::nullopt; - } - // The first lanes must map to registers in A - for (int i = 0; i < nRegA - nRegB; ++i) { - if (cvt.getBasis(kLane, laneIdTileB[i], kReg) == 0) { - return std::nullopt; - } - } - // The rest of the lanes must map to each other - for (auto [idxA, idxB] : - llvm::zip(laneIdTileA, laneIdTileB.take_back(laneIdTileA.size()))) { - if (cvt.getBasis(kLane, idxB, kLane) != (1 << idxA)) { - return std::nullopt; - } - } - vbasis.append(laneTileB.begin(), laneTileB.end()); - return vbasis; -} - -LinearLayout optimalSwizzling(const LinearLayout &src, const LinearLayout &dst, - int32_t bitwidth, ArrayRef vbasis, - ArrayRef tileSrc, - ArrayRef tileDst, - ArrayRef> outDims, - int32_t leaveReps = 0) { - // We work on the flattened tensors as the tensor dimensions are not relevant - assert(src.getNumOutDims() == 1 && dst.getNumOutDims() == 1 && - "src and dst must have a single output dimension"); - - const int32_t dim = src.getTotalOutDimSizeLog2(); - auto *ctx = src.getInDimNames().begin()->getContext(); - auto kReg = StringAttr::get(ctx, "register"); - - auto regsNotZero = [kReg](const LinearLayout &ll) { - return llvm::all_of( - ll.getBases().lookup(kReg), - [](const std::vector &basis) { return basis[0] != 0; }); - }; - assert( - regsNotZero(src) && - "Remove register broadcasting from src. See actionRemoveBroadcastedRegs"); - assert( - regsNotZero(dst) && - "Remove register broadcasting from dst. See actionRemoveBroadcastedRegs"); - - llvm::SmallVector bankSrc; - bankSrc.append(vbasis.begin(), vbasis.end()); - bankSrc.append(tileSrc.begin(), tileSrc.end()); - llvm::SmallVector bankDst; - bankDst.append(vbasis.begin(), vbasis.end()); - bankDst.append(tileDst.begin(), tileDst.end()); - - // Bits in a bank segment: 32 banks x 32 bits - constexpr int32_t bankBits = 32 * 32; - // Bases needed to cover a whole bank segment - const int32_t lenBbasis = std::min( - llvm::Log2_32(bankBits / ((1 << vbasis.size()) * bitwidth)), - dim - vbasis.size()); - // Bases to cover all the tensor - const int32_t lenSbasis = dim - lenBbasis - vbasis.size(); - - auto sbasis = computeSegment(bankSrc, bankDst, dim, lenSbasis); - - // The bank is the complement of the union of the vector and the start of the - // segments - SmallVector unionBasis; - unionBasis.append(vbasis.begin(), vbasis.end()); - unionBasis.append(sbasis.begin(), sbasis.end()); - SmallVector bbasis = complementBasis(unionBasis, dim); - - assert(bbasis.size() == lenBbasis + (lenSbasis - sbasis.size()) && - "bbasis size mismatch"); - - // Build the 1D result layout - StringAttr vecAttr = StringAttr::get(ctx, "vector"); - StringAttr bankAttr = StringAttr::get(ctx, "bank"); - StringAttr segAttr = StringAttr::get(ctx, "segment"); - - // src has just 1 outDim - LinearLayout basis1D({{vecAttr, unflatten(vbasis)}, - {bankAttr, unflatten(bbasis)}, - {segAttr, unflatten(sbasis)}}, - src.getOutDims(), /*requireSurjective=*/true); - basis1D = buildReps(ctx, src, dst, basis1D, leaveReps); - - return basis1D.reshapeOuts(outDims); -} -LinearLayout optimalSwizzlingLdSt(const LinearLayout &src, - const LinearLayout &dst, int32_t bitwidth) { - auto *ctx = src.getInDimNames().begin()->getContext(); - auto kReg = StringAttr::get(ctx, "register"); - auto kLane = StringAttr::get(ctx, "lane"); - auto srcFlat = src.flattenOuts(); - auto dstFlat = dst.flattenOuts(); - auto regSrc = flatten(srcFlat, kReg); - auto regDst = flatten(dstFlat, kReg); - auto laneSrc = flatten(srcFlat, kLane); - auto laneDst = flatten(dstFlat, kLane); - auto dim = src.getTotalOutDimSizeLog2(); - SmallVector vbasis = intersectionBasis(regSrc, regDst, dim); - // Restrict the vectorisation to the maximum we can use - auto maxVecBases = llvm::Log2_32(128 / bitwidth); - if (vbasis.size() > maxVecBases) { - vbasis.resize(maxVecBases); - } - // We fill-up vbasis until it has 32 bits as best we can - std::optional srcFillsBank = std::nullopt; - if ((1 << vbasis.size()) * bitwidth < 32) { - auto basesPerBank = llvm::Log2_32(32 / bitwidth); - auto kWarp = StringAttr::get(ctx, "warp"); - auto warpSrc = removeZeros(flatten(srcFlat, kWarp)); - auto warpDst = removeZeros(flatten(dstFlat, kWarp)); - auto removeVec = [&vbasis](ArrayRef vec) { - SmallVector result; - for (int32_t r : vec) { - if (!llvm::is_contained(vbasis, r)) { - result.push_back(r); - } - } - return result; - }; - auto regSrcWarp = intersectionBasis(removeVec(regSrc), warpDst, dim); - auto regDstWarp = intersectionBasis(removeVec(regDst), warpSrc, dim); - // Maximise vectorisation in the load or the store without creating - // conflicts - SmallVector largest; - if (regSrcWarp.size() == regDstWarp.size() && regSrcWarp.size() > 0) { - // We choose the one with the lowest basis in the hope that it will - // avoid PRMTs. The comparison of the mins will be strict as the sets - // removeVec(regSrc) and removeVec(regDst) don't intersect - if (*llvm::min_element(regSrcWarp) < *llvm::min_element(regDstWarp)) { - largest = regSrcWarp; - srcFillsBank = true; - } else { - largest = regDstWarp; - srcFillsBank = false; - } - } else { - srcFillsBank = regSrcWarp.size() > regDstWarp.size(); - largest = srcFillsBank.value() ? regSrcWarp : regDstWarp; - } - vbasis.append(largest.begin(), largest.end()); - - if (vbasis.size() < basesPerBank) { - // Pad the vectorisation to 32 bits with warp bases - auto warpSrcWarp = intersectionBasis(warpSrc, warpDst, dim); - vbasis.append(warpSrcWarp.begin(), warpSrcWarp.end()); - } - - int i = 0; - while (vbasis.size() < basesPerBank && - (i < warpSrc.size() || i < warpDst.size())) { - // If we have not filled up a whole bank, we add more warp bases - // until we have 32 bits. They will at least avoid bank conflicts in one - // direction - if (i < warpSrc.size() && !llvm::is_contained(vbasis, warpSrc[i])) { - vbasis.push_back(warpSrc[i]); - } - if (vbasis.size() < basesPerBank && i < warpDst.size() && - !llvm::is_contained(vbasis, warpDst[i])) { - vbasis.push_back(warpDst[i]); - } - ++i; - } - - // Trim to basesPerBank if we have added more - // The idea here is that implementing asymmetric vectorisation without bank - // conflicts is a bit tricky. Basically, in this case, you need to use the - // vectorisation base in the swizzling pattern. As such, you would not be - // able to vectorise all the `ld.shared` instructions that you emit, but - // just about half of them (the ones that are not swizzled). We don't - // implement this yet - if (vbasis.size() > basesPerBank) { - vbasis.resize(basesPerBank); - } - } - auto log2Vec = llvm::Log2_32( - std::max(1, ((1 << vbasis.size()) * bitwidth) / 32)); - auto tileSrc = to_vector(ArrayRef(laneSrc).drop_back(log2Vec)); - auto tileDst = to_vector(ArrayRef(laneDst).drop_back(log2Vec)); - auto smem = optimalSwizzling(srcFlat, dstFlat, bitwidth, vbasis, tileSrc, - tileDst, src.getOutDims()); - - // We might be able to vectorise a bit more the load or the store - // This may happen when there is broadcasting - // e.g for fp32 - // src = {reg = [], lane = [1, 2, 4, 8, 16], warp = [32]} - // dst = {reg = [8, 32], lane = [0, 0, 1, 2, 4], warp = [16]} - if (log2Vec < 2) { - auto smemFlat = smem.flattenOuts(); - // For every bank line, find if it is in regSrc or regDst - // and if so, store the index in the vector - SmallVector idxBanksInRegSrc; - SmallVector idxBanksInRegDst; - auto kBank = StringAttr::get(ctx, "bank"); - const auto &banks = flatten(smemFlat, kBank); - for (auto [i, r] : llvm::enumerate(banks)) { - if (llvm::is_contained(regSrc, r)) { - idxBanksInRegSrc.push_back(i); - } - if (llvm::is_contained(regDst, r)) { - idxBanksInRegDst.push_back(i); - } - } - - // Choose src/dst if we used them to fill the bank - // Otherwise choose the max vectorisation - SmallVector bBasisOrder; - if (srcFillsBank.has_value() && srcFillsBank.value()) { - bBasisOrder = std::move(idxBanksInRegSrc); - } else if (srcFillsBank.has_value() && !srcFillsBank.value()) { - bBasisOrder = std::move(idxBanksInRegDst); - } else { - bBasisOrder = idxBanksInRegSrc.size() > idxBanksInRegDst.size() - ? std::move(idxBanksInRegSrc) - : std::move(idxBanksInRegDst); - } - for (int i = 0; i < banks.size(); ++i) { - if (!llvm::is_contained(bBasisOrder, i)) { - bBasisOrder.push_back(i); - } - } - smem = ColumnAction(bBasisOrder, kBank, smem.getInDimSizeLog2(kBank)) - .apply(smem); - } - - return smem; -} - -std::pair> -optimalSwizzling(const LinearLayout &src, const LinearLayout &dst, - ArrayRef srcTiles, - ArrayRef dstTiles, int32_t bitwidth) { - assert(bitwidth <= 128 && "bitwidth must be <= 128"); - auto srcFlat = src.flattenOuts(); - auto dstFlat = dst.flattenOuts(); - // Number of total bases needed to cover the necessary contiguous tile - // We assume using ld.shared.b32.v4 in the case of ld/st ops - const auto totalBases = llvm::Log2_32(128 / bitwidth); - - auto *ctx = src.getInDimNames().begin()->getContext(); - auto kReg = StringAttr::get(ctx, "register"); - - // Find the pairs of instructions that we can use to lower this converet - SmallVector, SmallVector>> - instr; - for (const auto &[idxSrc, instrSrc] : llvm::enumerate(srcTiles)) { - auto logRegSrc = totalBases - instrSrc.laneContig.size(); - for (const auto &[idxDst, instrDst] : llvm::enumerate(dstTiles)) { - auto logRegDst = totalBases - instrDst.laneContig.size(); - auto maybeTile = - optimalSwizzlingTile(srcFlat, dstFlat, logRegSrc, logRegDst, - instrSrc.laneContig, instrDst.laneContig); - if (maybeTile.has_value()) { - instr.push_back({{idxSrc, idxDst}, std::move(*maybeTile)}); - } - } - } - auto getTile = - [](const LocalMemOpTile &instr, ArrayRef regs, - ArrayRef lane, - ArrayRef vbasis) -> std::optional> { - // pick the first 3 - laneAddr.size() registers that are not in vbasis - SmallVector tile; - auto regNeeded = 3 - instr.laneAddr.size(); - assert(regNeeded >= 0 && "laneAddr.size() must be <= 3"); - for (int32_t r : regs) { - if (regNeeded == 0) { - break; - } - if (!llvm::is_contained(vbasis, r)) { - tile.push_back(r); - regNeeded--; - } - } - // Not enough registers to fill in the tile - if (regNeeded > 0) { - return std::nullopt; - } - for (auto i : instr.laneAddr) { - tile.push_back(lane[i]); - } - return tile; - }; - - auto kLane = StringAttr::get(ctx, "lane"); - auto regSrc = flatten(srcFlat, kReg); - auto regDst = flatten(dstFlat, kReg); - auto laneSrc = flatten(srcFlat, kLane); - auto laneDst = flatten(dstFlat, kLane); - - // Get the associated src/dst tiles for each instruction if they exist - SmallVector, SmallVector, - SmallVector, SmallVector, int32_t>> - tiles; - for (auto [instrs, vbasis] : instr) { - auto maybeTileSrc = - getTile(srcTiles[instrs.first], regSrc, laneSrc, vbasis); - auto maybeTileDst = - getTile(dstTiles[instrs.second], regDst, laneDst, vbasis); - if (!maybeTileSrc.has_value() || !maybeTileDst.has_value()) { - continue; - } - // Regs bases missing to get full vectorisation - auto regsMissing = [](const LocalMemOpTile &instr) { - return instr.laneContig.size() + instr.laneAddr.size() - 3; - }; - // We leave 2 reps for combinations of ldmatrix/stmatrix instructions - // to be able to fully vectorise them - int32_t leaveReps = std::min(regsMissing(srcTiles[instrs.first]), - regsMissing(dstTiles[instrs.second])); - assert((leaveReps == 0 || leaveReps == 2) && "leaveReps must be 0 or 2"); - tiles.push_back({instrs, std::move(vbasis), std::move(*maybeTileSrc), - std::move(*maybeTileDst), leaveReps}); - } - - if (tiles.empty()) { - // We lower to an ld / st, but can't use LDS128/STS128 - auto smem = optimalSwizzlingLdSt(src, dst, bitwidth); - return {smem, {0, 0}}; - } else { - SmallVector>> - smems; - // We choose the pair of instructions that minimises the total bank - // conflicts - for (auto [instrs, vbasis, tileSrc, tileDst, leaveReps] : tiles) { - auto smem = optimalSwizzling(srcFlat, dstFlat, bitwidth, vbasis, tileSrc, - tileDst, src.getOutDims(), leaveReps); - auto [read, write] = bankConflicts(tileSrc, tileDst, smem); - smems.push_back({read + write, smem, {instrs.first, instrs.second}}); - } - // Current heuristic: Minimise total bank conflicts - // We break ties looking at the number of rounds we do to move the data - auto kReps = StringAttr::get(ctx, "reps"); - auto it = llvm::min_element(smems, [kReps](const auto &a, const auto &b) { - return std::get<0>(a) < std::get<0>(b) || - (std::get<0>(a) == std::get<0>(b) && - std::get<1>(a).getInDimSize(kReps) > - std::get<1>(b).getInDimSize(kReps)); - }); - return {std::get<1>(*it), std::get<2>(*it)}; - } -} - -} // namespace mlir::triton::gpu diff --git a/third_party/mthreads/lib/Tools/LayoutUtils.cpp b/third_party/mthreads/lib/Tools/LayoutUtils.cpp deleted file mode 100644 index 815bf6d4b3..0000000000 --- a/third_party/mthreads/lib/Tools/LayoutUtils.cpp +++ /dev/null @@ -1,582 +0,0 @@ -#include "triton/Tools/LayoutUtils.h" -#include "triton/Tools/GenericSwizzling.h" - -namespace mlir::triton { - -static bool checkSquareSublayout(const LinearLayout &ll, - ArrayRef dimNames, - function_ref checkBasis) { - // The empty layout is the identity - if (dimNames.size() == 0) { - return true; - } - // Check that the input-output sizes are the same - LinearLayout sl = ll.sublayout(dimNames, dimNames); - for (StringAttr dim : dimNames) { - if (ll.getInDimSize(dim) != ll.getOutDimSize(dim)) { - return false; - } - } - // Once the inputs and output dimensions are the same, we can just check - // that the basis for the single remaining dimension is the identity. - sl = sl.flattenIns().flattenOuts(); - const auto &inDimBases = sl.getBases().begin()->second; - for (auto [b, basis] : llvm::enumerate(inDimBases)) { - if (!checkBasis(b, basis[0])) { - return false; - } - } - return true; -} - -bool squareSublayoutIsIdentity(const LinearLayout &ll, - ArrayRef dimNames) { - return checkSquareSublayout( - ll, dimNames, [](int b, int32_t basis) { return basis == (1 << b); }); -} - -LinearLayout -ensureLayoutNotLargerThan(const LinearLayout &layout, - const llvm::SmallDenseMap &shape, - bool broadcastRegisters) { - assert(shape.size() == layout.getNumOutDims()); - if (shape.empty()) { - return layout; - } - MLIRContext *ctx = shape.begin()->first.getContext(); - - auto bases = layout.getBases(); - - auto kRegister = StringAttr::get(ctx, "register"); - std::set broadcastedDims; - - for (auto outDim : llvm::enumerate(layout.getOutDimNames())) { - auto outDimName = outDim.value(); - int32_t actualSize = layout.getOutDimSize(outDimName); - int32_t desiredSize = shape.lookup(outDimName); - if (actualSize <= desiredSize) { - continue; - } - assert(actualSize % desiredSize == 0); - // - std::vector> sortedBases; - for (auto [inDimName, basis] : bases) { - for (size_t basisIdx = 0; basisIdx < basis.size(); basisIdx++) { - auto outValue = basis[basisIdx][outDim.index()]; - if (outValue == 0) { - continue; - } - assert(llvm::isPowerOf2_32(outValue)); - sortedBases.emplace_back(inDimName, basisIdx, outValue); - } - } - // From the largest basis to the smallest. - llvm::sort(sortedBases, - [](auto a, auto b) { return std::get<2>(a) > std::get<2>(b); }); - for (auto [inDimName, basisIdx, outValue] : sortedBases) { - if (actualSize <= desiredSize) { - break; - } - if (!broadcastRegisters && inDimName == kRegister) { - broadcastedDims.insert(basisIdx); - } else { - bases[inDimName][basisIdx][outDim.index()] = 0; - } - actualSize >>= 1; - } - } - if (!broadcastRegisters) { - // Remove broadcasted registers - std::vector> newBasesRegister; - for (auto [idx, basis] : llvm::enumerate(bases[kRegister])) { - // Remove if it's broadcasted - if (broadcastedDims.find(idx) == broadcastedDims.end()) { - newBasesRegister.push_back(std::move(basis)); - } - } - bases[kRegister] = std::move(newBasesRegister); - } - auto outDims = layout.getOutDims(); - for (auto &[outDim, outDimSize] : outDims) { - outDimSize = std::min(outDimSize, shape.lookup(outDim)); - } - - return LinearLayout(std::move(bases), std::move(outDims), - /*requireSurjective=*/false); -} - -// For each out-dim d, ensure the layout's out-size (i.e. its codomain) is no -// smaller than shape[d]. Do this by increasing the size of the layout's inputs -// along its most-minor dimension ("register" for register layouts, "offset" for -// shared layouts). -// -// This function is invariant to the order of the layout's input dimensions, but -// it cares about the order of the output dims, which should be minor-to-major. -LinearLayout ensureLayoutNotSmallerThan( - const LinearLayout &layout, - const llvm::SmallDenseMap &shape) { - assert(shape.size() == layout.getNumOutDims()); - if (shape.empty()) { - return layout; - } - - StringAttr kDim = *layout.getInDimNames().begin(); - assert(kDim == "register" || kDim == "offset"); - - LinearLayout ret = layout; - for (StringAttr outDimName : layout.getOutDimNames()) { - int32_t actualSize = layout.getOutDimSize(outDimName); - int32_t desiredSize = shape.lookup(outDimName); - assert(actualSize > desiredSize || desiredSize % actualSize == 0); - ret *= LinearLayout::identity1D(desiredSize / actualSize, kDim, outDimName); - assert(ret.getOutDimSize(outDimName) >= desiredSize); - } - return ret; -} - -// Returns ["dim0", "dim1", ..., "dim"]. -SmallVector standardOutDimNames(MLIRContext *ctx, int rank) { - SmallVector ret; - for (int i = 0; i < rank; i++) { - ret.push_back(StringAttr::get(ctx, "dim" + llvm::Twine(i))); - } - return ret; -} - -// Returns [("dim0", dstShape[0]), ("dim1", dstShape[1]), ..., -// ("dim", dstShape[rank-1])]. -SmallVector> -standardOutDimPairs(MLIRContext *ctx, ArrayRef dstShape) { - auto newRank = dstShape.size(); - SmallVector> newOutDims; - for (auto [dim, size] : - llvm::zip(standardOutDimNames(ctx, newRank), dstShape)) { - newOutDims.emplace_back(dim, size); - } - return newOutDims; -} - -// Returns a 1D -> ND layout into [dim0, dim1, ...] that's equivalent to -// creating a 1D -> 1D mapping of size product(shape) and then reshaping to -// permute(shape, order). -LinearLayout identityStandardND(StringAttr inDimName, ArrayRef shape, - ArrayRef order) { - assert(shape.size() == order.size()); - MLIRContext *ctx = inDimName.getContext(); - auto rank = shape.size(); - - // The order in triton is written wrt. [dim0, dim1, ...]. - SmallVector outDimNames = standardOutDimNames(ctx, rank); - - LinearLayout ret = LinearLayout::empty(); - for (int i = 0; i < shape.size(); i++) { - // Start with the most-minor dimension, which is order[0]. - int dim = order[i]; - ret *= LinearLayout::identity1D(shape[dim], inDimName, outDimNames[dim]); - } - return ret; -} - -LinearLayout zerosLike(const LinearLayout &layout) { - auto bases = layout.getBases(); - for (auto &basis : bases) { - for (auto &vec : basis.second) { - for (auto &val : vec) { - val = 0; - } - } - } - - SmallVector> outDims; - for (auto outDim : layout.getOutDimNames()) { - outDims.emplace_back(outDim, layout.getOutDimSize(outDim)); - } - return LinearLayout(std::move(bases), std::move(outDims), - /*requireSurjective=*/false); -} - -std::optional regPermForDivide(const LinearLayout &A, - const LinearLayout &B, bool left) { - // We can implement this generically for any dimension, but for now we only do - // it for regs to keep the API simpler - assert(A.getNumInDims() != 0); - auto kReg = *A.getInDimNames().begin(); - assert(kReg.str() == "register"); - assert(B.getNumInDims() != 0); - assert(kReg == *B.getInDimNames().begin()); - - // We broadcast B to have the same number of out dims as A. - LinearLayout broadcast; - for (StringAttr out : A.getOutDimNames()) { - broadcast *= LinearLayout::identity1D(1, kReg, out); - } - auto BBroadcast = broadcast * B; - - // Retrieve the register bases from A and B. - const auto &ARegBases = A.getBases().lookup(kReg); - const auto &BRegBases = BBroadcast.getBases().lookup(kReg); - - llvm::DenseMap log2QuotSize; - for (StringAttr out : A.getOutDimNames()) { - log2QuotSize[out] = - A.getOutDimSizeLog2(out) - BBroadcast.getOutDimSizeLog2(out); - if (log2QuotSize[out] < 0) - return std::nullopt; - } - - auto multiplyByTileSize = - [&](ArrayRef bBasis) -> std::vector { - std::vector result; - size_t idx = 0; - assert(bBasis.size() == A.getNumOutDims()); - for (auto [dim, b] : llvm::zip(A.getOutDimNames(), bBasis)) { - result.push_back(b << log2QuotSize.lookup(dim)); - } - return result; - }; - - // Compute the permutation order: - // For each basis in B (in order), find its index in A (using each index at - // most once). We make sure we use each index at most once in case B - // broadcasts (weird case, but better safe than sorry). - SmallVector bIndices; - SmallVector used(ARegBases.size(), false); - for (auto bB : BRegBases) { - bool found = false; - if (!left) - bB = multiplyByTileSize(bB); - - for (size_t j = 0; j < ARegBases.size(); ++j) { - found = !used[j] && (ARegBases[j] == bB); - if (found) { - bIndices.push_back(j); - used[j] = true; - break; - } - } - if (!found) - return std::nullopt; // A basis from B not found in A. - } - // Append remaining indices from A (preserving their original order). - SmallVector remainingIndices; - for (size_t i = 0; i < ARegBases.size(); ++i) { - if (!used[i]) - remainingIndices.push_back(i); - } - SmallVector permOrder = to_vector(llvm::concat( - left ? bIndices : remainingIndices, left ? remainingIndices : bIndices)); - return ColumnAction(permOrder, kReg, ARegBases.size()); -} - -ColumnAction actionRemoveBroadcastedRegs(const LinearLayout &layout) { - assert(layout.getNumInDims() != 0); - auto kReg = *layout.getInDimNames().begin(); - assert(kReg.str() == "register"); - - // Drop the bases that are zero - const auto &bases = layout.getBases().lookup(kReg); - SmallVector permOrder; - for (size_t i = 0; i < bases.size(); ++i) { - if (!llvm::all_of(bases[i], [](size_t x) { return x == 0; })) { - permOrder.push_back(i); - } - } - return ColumnAction(permOrder, kReg, bases.size()); -} -std::pair -actionAdditiveStrides(const LinearLayout &layout, const LinearLayout addrLayout, - uint64_t maskSpanOffsets) { - // We are looking to put at the front (after any zeros) any basis that does - // not intersect with any bit moved by any basis in kLane / kWarp - // and that is not moved by any affine offset - - // Note this function assumes that if any registers are used in the addrLayout - // of the layout (as in ldmatrix/stmatrix) they will be the first non-zero - // registers within `layout` - assert(layout.getNumInDims() != 0); - auto kReg = *layout.getInDimNames().begin(); - assert(kReg.str() == "register"); - auto kLane = StringAttr::get(kReg.getContext(), "lane"); - auto kWarp = StringAttr::get(kReg.getContext(), "warp"); - assert(layout.getNumOutDims() == 1); - uint32_t bits = maskSpanOffsets; - llvm::SetVector tileBases; - for (auto bases : llvm::make_second_range(addrLayout.getBases())) { - for (auto basis : bases) { - bits |= basis[0]; - tileBases.insert(basis[0]); - } - } - SmallVector front, back; - for (auto [idx, basis] : llvm::enumerate(layout.getBases().lookup(kReg))) { - if ((basis[0] & bits) == 0 || tileBases.contains(basis[0])) { - front.push_back(idx); - } else { - back.push_back(idx); - } - } - auto permOrder = to_vector(llvm::concat(front, back)); - return {1 << front.size(), - ColumnAction(permOrder, kReg, layout.getInDimSizeLog2(kReg))}; -} - -SmallVector broadcastAs(const SmallVector &values, - const LinearLayout &layout) { - assert(layout.getNumInDims() != 0); - auto kReg = *layout.getInDimNames().begin(); - assert(kReg.str() == "register"); - uint32_t broadcastMask = layout.getFreeVariableMasks().lookup(kReg); - assert((layout.getInDimSize(kReg) / (1 << llvm::popcount(broadcastMask))) == - values.size()); - - std::vector> newBases; - int i = 0; - for (int j = 0; j < layout.getInDimSizeLog2(kReg); j++) { - if (broadcastMask & (1 << j)) { - newBases.push_back({0}); - } else { - newBases.push_back({1 << i}); - i++; - } - } - auto newLayout = LinearLayout({{kReg, std::move(newBases)}}, {kReg}); - SmallVector ret; - - ret.reserve(newLayout.getInDimSize(kReg)); - for (int i = 0; i < newLayout.getInDimSize(kReg); i++) { - int32_t srcIdx = newLayout.apply({{kReg, i}}).begin()->second; - ret.push_back(values[srcIdx]); - } - return ret; -} - -// Compute the supremum of two lists. -// If the supremum is not unique, we return the first list first -// Error out if the supremum does not exist -// e.g. sup([a, b], [a, c]) = [a, b, c], sup([a, b], [b, c]) = [a, b, c] -// sup([a, b], [b, a]) = error! Supremum does not exist. -SmallVector supremum(const SmallVector &x, - const SmallVector &y) { - llvm::SetVector result; - DenseMap posX, posY; - for (auto [idx, elem] : llvm::enumerate(x)) - posX[elem] = idx; - for (auto [idx, elem] : llvm::enumerate(y)) - posY[elem] = idx; - int i = 0, j = 0; - const int INF = std::numeric_limits::max(); - while (i < x.size() || j < y.size()) { - while (i < x.size() && result.contains(x[i])) - ++i; - while (j < y.size() && result.contains(y[j])) - ++j; - if (i >= x.size() && j >= y.size()) - break; - if (i < x.size() && j < y.size() && x[i] == y[j]) { - if (posY[x[i]] < j) - llvm_unreachable("Supremum does not exist"); - result.insert(x[i]); - ++i, ++j; - continue; - } - int candX = INF, candY = INF; - if (i < x.size()) { - if (posY.count(x[i]) && posY[x[i]] >= j) - candX = posY[x[i]]; - } - if (j < y.size()) { - if (posX.count(y[j]) && posX[y[j]] >= i) - candY = posX[y[j]]; - } - if (i < x.size() && candX == INF) { - result.insert(x[i]); - ++i; - continue; - } - if (j < y.size() && candY == INF) { - result.insert(y[j]); - ++j; - continue; - } - if (candX <= candY) { - if (posY[x[i]] < j) - llvm_unreachable("Supremum does not exist"); - result.insert(x[i]); - ++i; - } else { - if (posX[y[j]] < i) - llvm_unreachable("Supremum does not exist"); - result.insert(y[j]); - ++j; - } - } - return to_vector(result); -} - -LinearLayout reshapeLayout(MLIRContext *ctx, LinearLayout layout, - ArrayRef shape) { - int rank = shape.size(); - auto srcOutDims = to_vector(layout.getOutDimNames()); - std::reverse(srcOutDims.begin(), srcOutDims.end()); - auto newOutDims = standardOutDimPairs(ctx, shape); - std::reverse(newOutDims.begin(), newOutDims.end()); - return layout.transposeOuts(srcOutDims) - .reshapeOuts(newOutDims) - .transposeOuts(standardOutDimNames(ctx, rank)); -} - -LinearLayout transposeLinearLayout(LinearLayout layout, ArrayRef order) { - // Transpose the tile layout. - auto namedBases = layout.getBases(); - // move the most outer dimensions to the inner most position. - - for (auto &bases : llvm::make_second_range(namedBases)) { - for (auto &b : bases) { - std::vector newB; - for (auto i : order) { - newB.push_back(b[i]); - } - b = std::move(newB); - } - } - return LinearLayout(std::move(namedBases), - to_vector(layout.getOutDimNames())); -} - -std::pair -largestVectorisation(MLIRContext *ctx, const LinearLayout &cvt, int bitwidth, - std::optional maybeMaxVecElems) { - // Find the largest vectorisation we can use: - auto S = [ctx](StringRef str) { return StringAttr::get(ctx, str); }; - StringAttr kReg = S("register"); - StringAttr kOffset = S("offset"); - LinearLayout quot; - LinearLayout tile; - ColumnAction permutation; - // If there are restrictions on the vectorisation, we don't allow - // permutations. - auto allowPerm = !maybeMaxVecElems.has_value(); - auto maxVecElems = maybeMaxVecElems.value_or(128 / bitwidth); - for (int v = maxVecElems; v >= 1; v /= 2) { - tile = LinearLayout::identity1D(v, kReg, kOffset); - auto maybePerm = regPermForDivide(cvt, tile, /*left=*/true); - if (!maybePerm) { - continue; - } - permutation = *maybePerm; - if (!allowPerm && !permutation.isIdentity()) { - continue; - } - auto newCvt = permutation.apply(cvt); - auto maybeQuot = divideLeft(newCvt, tile); - if (!maybeQuot) { - continue; - } - return {v, permutation}; - } - llvm_unreachable("Vectorization < 1 is not valid"); -} - -std::optional getReps(const LinearLayout &cvt, - const LinearLayout &tile) { - - // Ensure tile out-dims are subset of cvt out-dims. - for (auto od : tile.getOutDimNames()) - assert(cvt.hasOutDim(od) && "tile out-dims must be contained in cvt"); - - // Precompute tile out-dim bit-widths. - llvm::SmallDenseMap outBLog2; - for (StringAttr od : cvt.getOutDimNames()) - outBLog2[od] = tile.hasOutDim(od) ? tile.getOutDimSizeLog2(od) : 0; - - // Build a per-out-dimension mask by OR-ing all tile bases that touch it. - llvm::SmallDenseMap tileMaskPerOutDim; - for (StringAttr od : cvt.getOutDimNames()) - tileMaskPerOutDim[od] = 0; - for (auto &[inDim, inBases] : tile.getBases()) { - (void)inDim; - for (auto &basis : inBases) { - int idx = 0; - for (StringAttr od : tile.getOutDimNames()) { - tileMaskPerOutDim[od] |= basis[idx++]; - } - } - } - - // Build reps with the same in/out dims as cvt, but zeroing out the leading - // inB bases (per in-dim) and keeping the remainder bases unchanged from cvt. - LinearLayout::BasesT repsBases; - for (StringAttr id : cvt.getInDimNames()) { - int inA = cvt.getInDimSizeLog2(id); - int inB = tile.hasInDim(id) ? tile.getInDimSizeLog2(id) : 0; - if (inB > inA) { - return std::nullopt; - } - - std::vector> basesForDim; - basesForDim.reserve(inA); - - // 1) Validate the starting bases match exactly. - for (int i = 0; i < inB; ++i) { - for (StringAttr od : cvt.getOutDimNames()) { - int a = cvt.getBasis(id, i, od); - int b = tile.getBasis(id, i, od); - if (a != b) { - return std::nullopt; - } - } - } - - // 2) Validate no overlap: the remaining cvt bases must have zeros in all - // tile-bit positions (computed as OR of all tile bases) for each - // out-dim. - for (int i = inB; i < inA; ++i) { - for (StringAttr od : cvt.getOutDimNames()) { - int32_t mask = tileMaskPerOutDim.lookup(od); - if (mask == 0) - continue; - int v = cvt.getBasis(id, i, od); - if ((v & mask) != 0) { - return std::nullopt; - } - } - } - - // 3) Emit reps bases: first inB as all-zeros; remainder copied from cvt. - for (int i = 0; i < inB; ++i) { - std::vector zero(cvt.getNumOutDims(), 0); - basesForDim.push_back(std::move(zero)); - } - for (int i = inB; i < inA; ++i) { - std::vector keep; - keep.reserve(cvt.getNumOutDims()); - for (StringAttr od : cvt.getOutDimNames()) - keep.push_back(cvt.getBasis(id, i, od)); - basesForDim.push_back(std::move(keep)); - } - - repsBases[id] = std::move(basesForDim); - } - - return LinearLayout(std::move(repsBases), cvt.getOutDims(), - /*requireSurjective=*/false); -} - -LinearLayout removeStandardDim(const LinearLayout &layout, int dim) { - auto rank = layout.getNumOutDims(); - assert(rank > 0); - assert(dim < rank); - auto *ctx = layout.getOutDimNames().begin()->getContext(); - auto dims = to_vector(layout.getOutDimNames()); - assert(dims == standardOutDimNames(ctx, rank)); - dims.erase(dims.begin() + dim); - auto newLayout = layout.sublayout(to_vector(layout.getInDimNames()), dims); - auto dimSizes = newLayout.getOutDims(); - auto newDims = standardOutDimNames(ctx, rank - 1); - for (auto [i, newDim] : llvm::enumerate(newDims)) { - dimSizes[i].first = newDim; - } - return LinearLayout(newLayout.getBases(), dimSizes, /*isSurjective*/ false); -} - -} // namespace mlir::triton diff --git a/third_party/mthreads/lib/Tools/LinearLayout.cpp b/third_party/mthreads/lib/Tools/LinearLayout.cpp deleted file mode 100644 index 11b4367072..0000000000 --- a/third_party/mthreads/lib/Tools/LinearLayout.cpp +++ /dev/null @@ -1,1407 +0,0 @@ -#include "triton/Tools/LinearLayout.h" - -#include -#include -#include - -#include "mlir/IR/BuiltinAttributes.h" -#include "third_party/f2reduce/f2reduce.h" -#include "triton/Tools/LayoutUtils.h" -#include "triton/Tools/StrUtil.h" -#include "llvm/ADT/STLExtras.h" -#include "llvm/ADT/SetOperations.h" -#include "llvm/ADT/StringRef.h" -#include "llvm/Support/Debug.h" -#include "llvm/Support/ErrorHandling.h" -#include "llvm/Support/MathExtras.h" - -#define DEBUG_TYPE "linear_layout" -#define DBGS() (llvm::dbgs() << "[" DEBUG_TYPE "]: ") -#define LDBG(X) LLVM_DEBUG(DBGS() << X << "\n") - -#if defined(_MSC_VER) && !defined(__clang__) -// from https://gist.github.com/pps83/3210a2f980fd02bb2ba2e5a1fc4a2ef0 -#include - -static int __builtin_ctz(unsigned x) { - unsigned long r; - _BitScanForward(&r, x); - return static_cast(r); -} - -static int __builtin_ctzll(unsigned long long x) { - unsigned long r; - _BitScanForward64(&r, x); - return static_cast(r); -} - -#endif - -namespace mlir::triton { - -namespace { -using BasesT = LinearLayout::BasesT; -using llvm::SmallDenseSet; -using llvm::Twine; - -BasesT makeBasesMap( - ArrayRef>>> bases) { - BasesT ret; - for (const auto &[inDim, inDimBases] : bases) { - ret[inDim] = inDimBases; - } - return ret; -} - -// Dump the matrix to stderr in a human-readable format for debugging. -void dumpMatrix(uint64_t *m, int numRows, int numCols) { - assert(numCols <= 64); - for (int r = 0; r < numRows; r++) { - llvm::errs() << "0b"; - for (int c = 0; c < numCols; c++) { - llvm::errs() << ((m[r] & (1 << c)) != 0 ? "1" : "0"); - } - llvm::errs() << "\n"; - } -} - -// Compute the rank of the matrix formed by taking the bases for the given -// outDim as columns. In other words, finds the number of linearly-independent -// bases for this output dimension. -int getMatrixRank(std::unique_ptr m, int numRows, int numCols) { - // stride is specified in number of 64-bit words per row, and we pack our - // matrix so that there's only one uint64_t per row. - assert(numCols <= 64); - f2reduce::inplace_rref_strided(m.get(), numRows, numCols, /*stride=*/1); - - // The rank of the reduced matrix is simply the number of nonzero rows. - int rank = 0; - for (int i = 0; i < numRows; i++) { - if (m[i] != 0) - rank++; - } - return rank; -} - -template -void assertDimsEqualIgnoringOrder(T &&a, U &&b) { - SmallDenseSet as(a.begin(), a.end()); - SmallDenseSet bs(b.begin(), b.end()); - if (as != bs) { - llvm::report_fatal_error("Dimensions must match, ignoring order, but they " - "don't. Got dims: [" + - Twine(triton::join(a, ", ")) + "] and [" + - triton::join(b, ", ") + "]"); - } -} - -template -void assertDimsSubsetIgnoringOrder(T &&small, U &&big) { - SmallDenseSet smallSet(small.begin(), small.end()); - SmallDenseSet bigSet(big.begin(), big.end()); - if (!llvm::set_is_subset(smallSet, bigSet)) { - llvm::report_fatal_error("Dimensions must be a subset, ignoring order, but " - "they aren't. Got dims: [" + - Twine(triton::join(small, ", ")) + "] and [" + - triton::join(big, ", ") + "]"); - } -} -} // anonymous namespace - -/*static*/ std::optional -LinearLayout::tryCreate(BasesT bases, - ArrayRef> outDims, - bool requireSurjective) { - LinearLayout ll(std::move(bases), std::move(outDims), NoCheckInvariants{}); - std::optional error = ll.checkInvariants(requireSurjective); - if (error) { - return std::nullopt; - } - return ll; -} - -LinearLayout::LinearLayout(BasesT bases, - ArrayRef> outDims, - NoCheckInvariants) - : bases(std::move(bases)) { - for (auto [outDim, size] : outDims) { - this->outDims[outDim] = size; - } -} - -LinearLayout::LinearLayout(BasesT bases, ArrayRef outDimNames) - : bases(std::move(bases)) { - // Infer out-dim sizes. - for (StringAttr outDim : outDimNames) { - outDims[outDim] = 1; - } - for (const auto &[inDim, inDimBases] : this->bases) { - for (const auto &basis : inDimBases) { - for (int i = 0; i < basis.size(); i++) { - int32_t &size = outDims[outDimNames[i]]; - size = std::max(size, llvm::NextPowerOf2(basis[i])); - } - } - } - - std::optional error = - checkInvariants(/*requireSurjective=*/true); - if (error.has_value()) { - llvm::report_fatal_error(StringRef(*error)); - } -} - -LinearLayout::LinearLayout(BasesT bases, - ArrayRef> outDims, - bool requireSurjective) - : LinearLayout(std::move(bases), std::move(outDims), NoCheckInvariants{}) { - std::optional error = checkInvariants(requireSurjective); - if (error.has_value()) { - llvm::report_fatal_error(StringRef(*error)); - } -} - -std::optional -LinearLayout::checkInvariants(bool requireSurjective) { - LDBG("checkInvariants: " << toString()); - // Check that basis values are non-negative. - for (const auto &[inDim, inDimBases] : bases) { - for (const auto &basis : inDimBases) { - if (llvm::any_of(basis, [](int32_t b) { return b < 0; })) { - return "Invalid bases passed to LinearLayout. Expected all basis " - "values to be non-negative, but found a negative value for " - "in dimension '" + - inDim.str() + "'. Full list of bases:" + toString() + "\n"; - } - } - } - - // Check that the bases all have length equal to outDimNames.size(). - for (const auto &[inDim, inDimBases] : bases) { - for (const auto &basis : inDimBases) { - if (basis.size() != outDims.size()) { - return "Invalid bases passed to LinearLayout. Expect all bases to " - "have the same size, equal to outDimNames.size() (" + - std::to_string(outDims.size()) + - "). But this failed for in dimension '" + inDim.str() + - "'. Full list of bases:" + toString() + "\n"; - } - } - } - - // Check that the out-dim sizes are powers of 2. - for (const auto &[outDim, size] : outDims) { - if (!llvm::isPowerOf2_32(size)) { - return "Invalid out-dim size " + std::to_string(size) + " for out-dim '" + - outDim.str() + "'. Out-dim sizes must be powers of 2.\n"; - } - } - - // Check that the bases are smaller than the out-dim sizes. - SmallVector outDimNames = llvm::to_vector(getOutDimNames()); - for (const auto &[inDim, inDimBases] : this->bases) { - for (const auto &basis : inDimBases) { - for (int i = 0; i < basis.size(); i++) { - if (basis[i] >= outDims[outDimNames[i]]) { - return "Invalid basis " + std::to_string(basis[i]) + " for in-dim '" + - inDim.str() + "' and out-dim '" + outDimNames[i].str() + - "'. Basis must be less than the out-dim size.\n"; - } - } - } - } - - // Determine whether the this layout is surjective, i.e. that every `out` - // coordinate can be reached by some `in` coordinate. - // - // It's prohibitively slow to calculate this naively, but thankfully, this - // is equivalent to checking that the number of linearly-independent bases - // is equal to sum(getOutDimSizeLog2). This can be computed by finding - // the rank of the matrix whose columns are those bases. We can compute - // the rank of our matrix using Gaussian elimination, which runs in O(n^3) - // for an n x n matrix. Our matrix size is sum(inDimSizeLog2) x - // sum(outDimSizeLog2), so this should be plenty fast. - this->rank = - getMatrixRank(getMatrix(*this), /*numRows=*/getTotalOutDimSizeLog2(), - /*numCols=*/getTotalInDimSizeLog2()); - - if (requireSurjective && !isSurjective()) { - return "Layout is expected to be surjective, i.e. every `out` coordinate " - "can be reached by some `in` coordinate, but was not:" + - toString(); - } - - return std::nullopt; -} - -LinearLayout::LinearLayout( - ArrayRef>>> bases, - ArrayRef outDimNames) - : LinearLayout(makeBasesMap(bases), outDimNames) {} - -LinearLayout::LinearLayout( - ArrayRef>>> bases, - ArrayRef> outDims, bool requireSurjective) - : LinearLayout(makeBasesMap(bases), outDims, requireSurjective) {} - -/*static*/ LinearLayout LinearLayout::strided1D(int32_t size, int32_t stride, - StringAttr inDimName, - StringAttr outDimName) { - if (size == 0) - return LinearLayout::empty(); - - assert(llvm::isPowerOf2_32(size)); - std::vector> bases; - for (int32_t i = 1; i < size; i *= 2) { - bases.emplace_back(std::vector{i * stride}); - } - bool requiresSurjective = (stride == 1); - return LinearLayout({{inDimName, std::move(bases)}}, - {{outDimName, stride * size}}, requiresSurjective); -} - -/*static*/ LinearLayout LinearLayout::zeros1D(int32_t size, - StringAttr inDimName, - StringAttr outDimName, - int32_t outDimSize) { - if (size == 0) - return LinearLayout::empty(); - - assert(llvm::isPowerOf2_32(size)); - std::vector> zeros; - for (int i = 1; i < size; i *= 2) { - zeros.emplace_back(std::vector{0}); - } - return LinearLayout({{inDimName, zeros}}, {{outDimName, outDimSize}}, - /*requiresSurjective=*/outDimSize == 1); -} - -int32_t LinearLayout::getOutDimIndex(StringAttr outDim) const { - int i = 0; - for (auto [name, _] : outDims) { - if (name == outDim) { - return i; - } - i++; - } - llvm::report_fatal_error("outDim " + Twine(outDim) + " is not in layout" + - toString()); -} - -int32_t LinearLayout::getInDimSizeLog2(StringAttr inDim) const { - auto it = bases.find(inDim); - assert(it != bases.end() && "inDim not found in layout"); - return it->second.size(); -} - -int32_t LinearLayout::getTotalInDimSizeLog2() const { - return std::accumulate(getInDimNames().begin(), getInDimNames().end(), 0, - [&](int32_t acc, StringAttr inDim) { - return acc + getInDimSizeLog2(inDim); - }); -} - -int32_t LinearLayout::getOutDimSizeLog2(StringAttr outDim) const { - auto it = outDims.find(outDim); - assert(it != outDims.end() && "outDim not found in layout"); - return llvm::Log2_32(it->second); -} - -int32_t LinearLayout::getTotalOutDimSizeLog2() const { - return std::accumulate(getOutDimNames().begin(), getOutDimNames().end(), 0, - [&](int32_t acc, StringAttr outDim) { - return acc + getOutDimSizeLog2(outDim); - }); -} - -int32_t LinearLayout::getNumConsecutiveInOut() const { - if (bases.empty() || getNumOutDims() == 0) - return 1; - - // Count how many of the initial bases for the first in-dim are - // (2^i, 0, ..., 0). - const auto &firstInDimBases = bases.begin()->second; - int consec = 0; - for (; consec < firstInDimBases.size(); consec++) { - const auto &basis = firstInDimBases[consec]; - if (basis[0] != (1 << consec) || - !std::all_of(basis.begin() + 1, basis.end(), - [](int32_t x) { return x == 0; })) { - break; - } - } - - // `or` together all other bases' first out-dim. - int32_t otherBits = 0; - for (const auto &[inDim, inDimBases] : bases) { - for (int i = 0; i < inDimBases.size(); i++) { - if (inDim != bases.begin()->first || i >= consec) { - otherBits |= inDimBases[i][0]; - } - } - } - int32_t trailingZeros = otherBits != 0 ? __builtin_ctz(otherBits) : 31; - - return 1 << std::min(consec, trailingZeros); -} - -LinearLayout LinearLayout::transposeIns(ArrayRef newInDims) const { - assertDimsEqualIgnoringOrder(newInDims, getInDimNames()); - - BasesT newBases; - for (const auto &inDim : newInDims) { - newBases[inDim] = bases.find(inDim)->second; - } - return LinearLayout(std::move(newBases), llvm::to_vector(outDims), - isSurjective()); -} - -LinearLayout -LinearLayout::transposeOuts(ArrayRef newOutDims) const { - assertDimsEqualIgnoringOrder(newOutDims, getOutDimNames()); - - std::vector permutation; - for (const auto &outDim : newOutDims) { - permutation.push_back(getOutDimIndex(outDim)); - } - - BasesT newBases; - for (const auto &[inDim, inDimBases] : bases) { - auto &newInDimBases = newBases[inDim]; - for (const auto &basis : inDimBases) { - std::vector newBasis; - for (int32_t i : permutation) { - newBasis.push_back(basis[i]); - } - newInDimBases.push_back(std::move(newBasis)); - } - } - - SmallVector> newOutDimSizes; - for (auto outDim : newOutDims) { - newOutDimSizes.push_back({outDim, getOutDimSize(outDim)}); - } - return LinearLayout(std::move(newBases), newOutDimSizes, isSurjective()); -} - -LinearLayout LinearLayout::reshapeIns( - ArrayRef> newInDims) const { - assert(llvm::all_of(newInDims, [&](auto &inDim) { - return llvm::isPowerOf2_32(inDim.second); - })); - assert(getTotalInDimSize() == std::accumulate(newInDims.begin(), - newInDims.end(), 1, - [&](int32_t acc, auto &inDim) { - return acc * inDim.second; - })); - - // First flatten into a single in-dimension. Then split it up according - // to `newInDims`. - SmallVector> flatBases; - for (const auto &[inDim, inDimBases] : bases) { - for (const auto &basis : inDimBases) { - flatBases.push_back(basis); - } - } - - BasesT newBases; - int i = 0; - for (const auto &[inDim, inDimSize] : newInDims) { - auto &newInDimBases = newBases[inDim]; - for (int j = 1; j < inDimSize; j *= 2) { - newInDimBases.push_back(flatBases[i++]); - } - } - return LinearLayout(std::move(newBases), llvm::to_vector(outDims), - isSurjective()); -} - -LinearLayout LinearLayout::reshapeOuts( - ArrayRef> newOutDims) const { - assert(llvm::all_of(newOutDims, [&](auto &outDim) { - return llvm::isPowerOf2_32(outDim.second); - })); - assert(getTotalOutDimSize() == - std::accumulate( - newOutDims.begin(), newOutDims.end(), 1, - [&](int32_t acc, auto &outDim) { return acc * outDim.second; })); - - SmallVector shifts; - shifts.push_back(0); - for (StringAttr outDim : getOutDimNames()) { - shifts.push_back(shifts.back() + getOutDimSizeLog2(outDim)); - } - - // Flatten into a single out-dimension. Then split it up according to - // `newOutDims`. - llvm::MapVector> flatBases; - for (const auto &[inDim, inDimBases] : bases) { - auto &flatInBases = flatBases[inDim]; - for (const auto &basis : inDimBases) { - int b = 0; - for (int i = 0; i < basis.size(); i++) { - b += basis[i] << shifts[i]; - } - flatInBases.push_back(b); - } - } - - BasesT newBases; - for (const auto &[inDim, flatInBases] : flatBases) { - std::vector> &newInDimBases = newBases[inDim]; - for (int32_t b : flatInBases) { - std::vector multiDimBasis; - for (int32_t newSize : llvm::make_second_range(newOutDims)) { - multiDimBasis.push_back(b % newSize); - b /= newSize; - } - newInDimBases.push_back(std::move(multiDimBasis)); - } - } - - return LinearLayout(std::move(newBases), newOutDims, isSurjective()); -} - -LinearLayout LinearLayout::resizeInDim(StringAttr inDim, - int32_t newSize) const { - assert(llvm::isPowerOf2_32(newSize)); - assert(newSize <= getInDimSize(inDim)); - auto newBases = bases; - newBases[inDim].resize(llvm::Log2_32(newSize)); - return LinearLayout(std::move(newBases), getOutDims(), - /*requiresSurjective=*/false); -} - -LinearLayout LinearLayout::resizeOutDim(StringAttr outDim, - int32_t newSize) const { - assert(llvm::isPowerOf2_32(newSize)); - assert(newSize <= getOutDimSize(outDim)); - auto newBases = bases; - // Zero-out the basis vectors that are greater than or equal to the new size - for (auto &[inDim, inDimBases] : newBases) { - for (auto &basis : inDimBases) { - auto &b = basis[getOutDimIndex(outDim)]; - if (b >= newSize) { - b = 0; - } - } - } - auto outDims = getOutDims(); - for (auto &[outDim, outDimSize] : outDims) { - if (outDim == outDim) { - outDimSize = newSize; - } - } - return LinearLayout(std::move(newBases), outDims, - /*requiresSurjective=*/false); -} - -LinearLayout LinearLayout::concatIns(const LinearLayout &other) const { - assert(llvm::to_vector(getOutDimNames()) == - llvm::to_vector(other.getOutDimNames()) && - "layouts must have the same output dimensions"); - for (StringAttr outDim : getOutDimNames()) { - assert(getOutDimSize(outDim) == other.getOutDimSize(outDim) && - "layouts must have the same output dimension sizes"); - } - - LinearLayout::BasesT resultBases = getBases(); - for (auto &bases : other.getBases()) - resultBases.insert(bases); - SmallVector> newOutDims; - for (auto &[outDim, outDimSize] : outDims) - newOutDims.emplace_back(outDim, outDimSize); - return LinearLayout(std::move(resultBases), newOutDims, - /*requiresSurjective=*/false); -} - -LinearLayout LinearLayout::concatOuts(const LinearLayout &other) const { - assert(llvm::to_vector(getInDimNames()) == - llvm::to_vector(other.getInDimNames()) && - "layouts must have the same input dimensions"); - for (StringAttr inDim : getInDimNames()) { - assert(getInDimSize(inDim) == other.getInDimSize(inDim) && - "layouts must have the same input dimension sizes"); - } - - LinearLayout::BasesT result; - for (auto [lhsBases, rhsBases] : llvm::zip(getBases(), other.getBases())) { - auto &resultBases = result[lhsBases.first]; - assert(lhsBases.first == rhsBases.first); - for (auto [lhsBasis, rhsBasis] : - llvm::zip(lhsBases.second, rhsBases.second)) { - std::vector resultBasis; - llvm::append_range(resultBasis, lhsBasis); - llvm::append_range(resultBasis, rhsBasis); - resultBases.push_back(std::move(resultBasis)); - } - } - SmallVector> newOutDims; - for (auto &[outDim, outDimSize] : outDims) - newOutDims.emplace_back(outDim, outDimSize); - for (auto &[outDim, outDimSize] : other.outDims) - newOutDims.emplace_back(outDim, outDimSize); - return LinearLayout(std::move(result), newOutDims, - /*requiresSurjective=*/false); -} - -std::optional divideLeft(const LinearLayout &A, - const LinearLayout &B) { - // Compute a C such that A = B * C if it exists. - // Note that such a C exists iff (every pair of input/output dim of) A is of - // the form - // [[B, 0], - // [0, C]] - // as a matrix, whenever those dimensions are present in B. - for (StringAttr dim : B.getInDimNames()) { - if (!llvm::is_contained(A.getInDimNames(), dim)) - return std::nullopt; - } - for (StringAttr dim : B.getOutDimNames()) { - if (!llvm::is_contained(A.getOutDimNames(), dim)) - return std::nullopt; - } - // Compute candidate C's log-sizes for output dimensions. - llvm::MapVector cOutDimSizes; - for (StringAttr outDim : A.getOutDimNames()) { - int outA = A.getOutDimSizeLog2(outDim); - int outB = B.hasOutDim(outDim) ? B.getOutDimSizeLog2(outDim) : 0; - int outC = outA - outB; - if (outC < 0) - return std::nullopt; - cOutDimSizes[outDim] = 1 << outC; - } - - LinearLayout::BasesT cBases; - for (StringAttr inDim : A.getInDimNames()) { - int inA = A.getInDimSizeLog2(inDim); - int inB = B.hasInDim(inDim) ? B.getInDimSizeLog2(inDim) : 0; - int inC = inA - inB; - if (inC < 0) - return std::nullopt; - - std::vector> basesForDim; - // Check that A’s first inB entries agree with B. - for (int i = 0; i < inB; ++i) { - for (StringAttr outDim : A.getOutDimNames()) { - int expected = B.hasOutDim(outDim) ? B.getBasis(inDim, i, outDim) : 0; - int actual = A.getBasis(inDim, i, outDim); - if (actual != expected) - return std::nullopt; - } - } - - // Extract the candidate C bases from the remaining (shifted) entries in A. - for (int i = inB; i < inA; ++i) { - std::vector candidateBasis; - for (StringAttr outDim : llvm::make_first_range(cOutDimSizes)) { - int outB = B.hasOutDim(outDim) ? B.getOutDimSizeLog2(outDim) : 0; - int v = A.getBasis(inDim, i, outDim); - - // The lower outB bits must be zero. - if ((v & ((1 << outB) - 1)) != 0) - return std::nullopt; - candidateBasis.push_back(v >> outB); - } - basesForDim.push_back(std::move(candidateBasis)); - } - cBases[inDim] = basesForDim; - } - - SmallVector> COutDims; - for (auto [outDim, outC] : cOutDimSizes) { - COutDims.push_back({outDim, outC}); - } - // If the layout A and B are surjective, then C should also be surjective. - LinearLayout C(std::move(cBases), COutDims, - /*requireSurjective=*/A.isSurjective() && B.isSurjective()); - assert(B * C == A); - return C; -} - -std::optional divideRight(const LinearLayout &A, - const LinearLayout &B) { - // Compute a C such that A = C * B if it exists. - // Note that such a C exists iff (every pair of input/output dim of) A is of - // the form - // [[C, 0], - // [0, B]] - // as a matrix, whenever those dimensions are present in B. - - // Check that B's in-dimensions and out-dimensions are contained in A. - for (StringAttr dim : B.getInDimNames()) { - if (!llvm::is_contained(A.getInDimNames(), dim)) - return std::nullopt; - } - for (StringAttr dim : B.getOutDimNames()) { - if (!llvm::is_contained(A.getOutDimNames(), dim)) - return std::nullopt; - } - - // Compute candidate C's log-sizes for output dimensions. - llvm::MapVector cOutDimSizes; - for (StringAttr outDim : A.getOutDimNames()) { - int outA = A.getOutDimSizeLog2(outDim); - int outB = B.hasOutDim(outDim) ? B.getOutDimSizeLog2(outDim) : 0; - int outC = outA - outB; - if (outC < 0) - return std::nullopt; - cOutDimSizes[outDim] = 1 << outC; - } - - // For candidate C, its in-dim sizes come from subtracting B's in-dim sizes - // from A's. - LinearLayout::BasesT cBases; - for (StringAttr inDim : A.getInDimNames()) { - int inA = A.getInDimSizeLog2(inDim); - int inB = B.hasInDim(inDim) ? B.getInDimSizeLog2(inDim) : 0; - int inC = inA - inB; - if (inC < 0) - return std::nullopt; - - std::vector> basesForDim; - // The first inC basis vectors come directly from C. - for (int i = 0; i < inC; ++i) { - std::vector candidate; - for (StringAttr outDim : llvm::make_first_range(cOutDimSizes)) { - candidate.push_back(A.getBasis(inDim, i, outDim)); - } - basesForDim.push_back(std::move(candidate)); - } - - // The remaining inB basis vectors in A should correspond to B after being - // shifted. - for (int i = inC; i < inA; ++i) { - int j = i - inC; // Index into B's basis vectors for this inDim. - for (StringAttr outDim : B.getOutDimNames()) { - int outA = A.getOutDimSizeLog2(outDim); - int outB = B.getOutDimSizeLog2(outDim); - int outC = outA - outB; // Expected log2 size for C in this output. - int shift = outC; - int v = A.getBasis(inDim, i, outDim); - // The lower shift bits must be zero. - if ((v & ((1 << shift) - 1)) != 0) - return std::nullopt; - int recovered = v >> shift; - int expected = B.getBasis(inDim, j, outDim); - if (recovered != expected) - return std::nullopt; - } - } - cBases[inDim] = basesForDim; - } - - SmallVector> COutDims; - for (auto [outDim, size] : cOutDimSizes) - COutDims.push_back({outDim, size}); - // If A and B are surjective, then C should also be surjective. - LinearLayout C(std::move(cBases), COutDims, - /*requireSurjective=*/A.isSurjective() && B.isSurjective()); - assert(C * B == A); - return C; -} - -LinearLayout operator*(LinearLayout inner, LinearLayout outer) { - // Check that dims common to outer and inner have the same relative order. - auto inDims = supremum(llvm::to_vector(inner.getInDimNames()), - llvm::to_vector(outer.getInDimNames())); - auto outDims = supremum(llvm::to_vector(inner.getOutDimNames()), - llvm::to_vector(outer.getOutDimNames())); - - // Get the sizeLog2 of all input and output dimensions we're going to - // consider, in order. `inner` is more minor, so its dimensions come - // first. - llvm::MapVector inDimSizesLog2; - llvm::MapVector outDimSizesLog2; - for (const auto &dim : inDims) - inDimSizesLog2.insert({dim, 0}); - for (const auto &dim : outDims) - outDimSizesLog2.insert({dim, 0}); - for (const auto &layout : {inner, outer}) { - for (StringAttr inDim : layout.getInDimNames()) { - inDimSizesLog2[inDim] += layout.getInDimSizeLog2(inDim); - } - for (StringAttr outDim : layout.getOutDimNames()) { - outDimSizesLog2[outDim] += layout.getOutDimSizeLog2(outDim); - } - } - - BasesT allBases; - for (auto [inDimName, inDimSizeLog2] : inDimSizesLog2) { - std::vector> &inDimBases = allBases[inDimName]; - - // Fill with zeros. - inDimBases = std::vector>( - inDimSizeLog2, std::vector(outDimSizesLog2.size(), 0)); - - for (auto [outDimIdx, outDimNameAndSize] : - llvm::enumerate(outDimSizesLog2)) { - auto [outDimName, outDimSize] = outDimNameAndSize; - if (inner.hasInDim(inDimName) && inner.hasOutDim(outDimName)) { - for (int i = 0; i < inner.getInDimSizeLog2(inDimName); i++) { - inDimBases[i][outDimIdx] = inner.getBasis(inDimName, i, outDimName); - } - } - if (outer.hasInDim(inDimName) && outer.hasOutDim(outDimName)) { - int offset = - inner.hasInDim(inDimName) ? inner.getInDimSizeLog2(inDimName) : 0; - int shift = inner.hasOutDim(outDimName) - ? inner.getOutDimSizeLog2(outDimName) - : 0; - for (int i = 0; i < outer.getInDimSizeLog2(inDimName); i++) { - inDimBases[offset + i][outDimIdx] = - outer.getBasis(inDimName, i, outDimName) << shift; - } - } - } - } - - llvm::SmallVector> outDimSizes; - for (auto [outDim, sizeLog2] : outDimSizesLog2) { - outDimSizes.push_back({outDim, 1 << sizeLog2}); - } - return LinearLayout(std::move(allBases), outDimSizes, - inner.isSurjective() && outer.isSurjective()); -} - -bool LinearLayout::isTrivialOver(ArrayRef dimNames) const { - for (StringAttr dim : dimNames) { - if (!llvm::is_contained(getInDimNames(), dim) && - !llvm::is_contained(getOutDimNames(), dim)) { - return false; - } - } - - auto getRemainingDimNames = [&](auto allDimNames) { - SmallVector remainingDimNames; - for (StringAttr dim : allDimNames) { - if (!llvm::is_contained(dimNames, dim)) { - remainingDimNames.push_back(dim); - } - } - return remainingDimNames; - }; - SmallVector remainingInDimNames = - getRemainingDimNames(getInDimNames()); - SmallVector remainingOutDimNames = - getRemainingDimNames(getOutDimNames()); - - // Think of this as a block-matrix multiplying a vector: - // [[A, B], * [v_1, - // [C, D]] v_2] - // where v_2 is the dimNames and v_1 is the remainingInDimNames - // We can quotient out dimNames iff they don't affect the remainingInDimNames - // in the result. In other words, we want to check that B is zero, and C is - // zero, and D is the identity - return squareSublayoutIsIdentity(*this, dimNames) && - sublayoutIsZero(remainingInDimNames, dimNames) && - sublayoutIsZero(dimNames, remainingOutDimNames); -} - -std::optional -LinearLayout::quotient(ArrayRef dimNames) const { - if (!isTrivialOver(dimNames)) { - return std::nullopt; - } - - // This should probably be even less general, where we ask inDimNames == - // outDimNames - auto getRemainingDimNames = [&](auto allDimNames) { - SmallVector remainingDimNames; - for (StringAttr dim : allDimNames) { - if (!llvm::is_contained(dimNames, dim)) { - remainingDimNames.push_back(dim); - } - } - return remainingDimNames; - }; - - SmallVector inDimNames = getRemainingDimNames(getInDimNames()); - SmallVector outDimNames = getRemainingDimNames(getOutDimNames()); - - return sublayout(inDimNames, outDimNames); -} - -LinearLayout LinearLayout::sublayout(ArrayRef inDimNames, - ArrayRef outDimNames) const { - assertDimsSubsetIgnoringOrder(inDimNames, getInDimNames()); - assertDimsSubsetIgnoringOrder(outDimNames, getOutDimNames()); - SmallDenseSet inDimSet(inDimNames.begin(), inDimNames.end()); - SmallDenseSet outDimSet(outDimNames.begin(), outDimNames.end()); - - SmallVector outDimIndicesToKeep; - for (auto [i, outDim] : llvm::enumerate(getOutDimNames())) { - if (outDimSet.contains(outDim)) { - outDimIndicesToKeep.push_back(i); - } - } - BasesT newBases; - for (auto [inDim, inDimBases] : bases) { - if (!inDimSet.contains(inDim)) { - continue; - } - auto &newInDimBases = newBases[inDim]; - for (auto &basis : inDimBases) { - auto &newBasis = newInDimBases.emplace_back(); - for (int i : outDimIndicesToKeep) { - newBasis.push_back(basis[i]); - } - } - } - - SmallVector> newOutDims; - for (auto [outDim, outDimSize] : outDims) { - if (outDimSet.contains(outDim)) { - newOutDims.push_back({outDim, outDimSize}); - } - } - return LinearLayout(std::move(newBases), std::move(newOutDims), - /*requireSurjective=*/false); -} - -bool LinearLayout::sublayoutIsZero(ArrayRef inDimNames, - ArrayRef outDimNames) const { - LinearLayout ss = sublayout(inDimNames, outDimNames); - for (auto [inDim, inDimBases] : ss.bases) { - for (auto basis : inDimBases) { - if (!llvm::all_of(basis, [](int32_t b) { return b == 0; })) { - return false; - } - } - } - return true; -} - -SmallVector> -LinearLayout::apply(ArrayRef> ins) const { - assertDimsEqualIgnoringOrder(llvm::make_first_range(ins), getInDimNames()); - - SmallVector> ret; - for (StringAttr outDim : getOutDimNames()) { - int32_t outVal = 0; - for (auto &[inDim, val] : ins) { - for (int i = 0; i < getInDimSizeLog2(inDim); i++) { - if (val & (1 << i)) - outVal ^= getBasis(inDim, i, outDim); - } - } - ret.push_back({outDim, outVal}); - } - return ret; -} - -LinearLayout LinearLayout::compose(const LinearLayout &outer) const { - assertDimsEqualIgnoringOrder(getOutDimNames(), outer.getInDimNames()); - for (StringAttr outDim : getOutDimNames()) { - assert(getOutDimSize(outDim) <= outer.getInDimSize(outDim)); - } - - BasesT newBases; - for (const auto &[inDim, inDimBases] : bases) { - auto &newInDimBases = newBases[inDim]; - for (const auto &basis : inDimBases) { - SmallVector> bases; - for (auto [outDim, b] : llvm::zip(getOutDimNames(), basis)) { - bases.push_back({outDim, b}); - } - auto newBases = outer.apply(bases); - auto newBasesRange = llvm::make_second_range(newBases); - newInDimBases.push_back( - std::vector(newBasesRange.begin(), newBasesRange.end())); - } - } - - bool compositionIsSurjective = - isSurjective() && outer.isSurjective() && - llvm::all_of(getOutDimNames(), [&](StringAttr outDim) { - return getOutDimSize(outDim) == outer.getInDimSize(outDim); - }); - return LinearLayout(std::move(newBases), llvm::to_vector(outer.outDims), - compositionIsSurjective); -} - -namespace { -std::unique_ptr concatMatrices(const LinearLayout &A, - const LinearLayout &B) { - // conv - assert(A.getTotalOutDimSizeLog2() >= B.getTotalOutDimSizeLog2() && - "A must have at least as many output bits as B"); - int numColsA = A.getTotalInDimSizeLog2(); - - // rref expects the lower bits to be the lower indices of the matrix - auto concat = getMatrix(A); - auto BMat = getMatrix(B); - int rowA = 0; - int rowB = 0; - for (auto [outDim, outDimSize] : A.getOutDims()) { - for (int r = 0; r < llvm::Log2_32(outDimSize); r++) { - if (r < llvm::Log2_32(B.getOutDimSize(outDim))) { - concat[rowA] |= BMat[rowB] << numColsA; - rowB++; - } - rowA++; - } - } - return concat; -} - -LinearLayout lstsq(const LinearLayout &A, const LinearLayout &B) { - // Solve the least square system AX = B - // and return the least square solution X by computing RREF and setting - // the free variables to zero. - // A and B may not be surjective, but we assume that Im(B) \subset Im(A) - // Sketch of the algorithm: - // https://github.com/triton-lang/triton/pull/5309#discussion_r1869084111 - int numRows = A.getTotalOutDimSizeLog2(); - assert(numRows >= B.getTotalOutDimSizeLog2() && - "A.lstsq(B) called with incompatible output shapes"); - int numColsA = A.getTotalInDimSizeLog2(); - int numColsB = B.getTotalInDimSizeLog2(); - int numCols = numColsA + numColsB; - std::unique_ptr combinedMat = concatMatrices(A, B); - f2reduce::inplace_rref_strided(combinedMat.get(), numRows, numCols, - /*stride=*/1); - - // Compute the pivot columns - // Since A and B have the same image, each row will either have a pivot - // or will be all zeros - SmallVector pivotRowOfCol(numColsA, -1); - for (int r = 0; r < numRows; r++) { - auto row = combinedMat[r]; - if (row == 0) { - continue; - } - int c = __builtin_ctzll(row); - assert(c < numColsA && "Precondition broken. Im(B) not contained in Im(A)"); - assert(pivotRowOfCol[c] == -1 && - "duplicate pivot => matrix not in RREF or A not injective"); - pivotRowOfCol[c] = r; - } - - // Extract A^{-1}B and complete the matrix using zeros - std::unique_ptr retMat(new uint64_t[numColsA]()); - for (int c = 0; c < numColsA; ++c) { - int row = pivotRowOfCol[c]; - retMat[c] = (row == -1) ? 0 : (combinedMat[row] >> numColsA); - } - - // We need names for the in/out dim of the flattened layout we're going to - // read off from `m`. These could be anything, doesn't matter. - assert(!A.getInDimNames().empty() && - "attempt to solve lstsq for empty layout"); - StringAttr inDim1D = *A.getInDimNames().begin(); - StringAttr outDim1D = *A.getOutDimNames().begin(); - - // Read off the new bases. These are for a flattened 1D -> 1D - LinearLayout::BasesT retBases; - auto &bs = retBases[inDim1D]; - for (int c = 0; c < numColsB; c++) { - int32_t basis = 0; - for (int r = 0; r < numColsA; r++) { - basis |= (retMat[r] >> c & 1) << r; - } - bs.push_back({basis}); - } - - LinearLayout retFlattened(std::move(retBases), - {{outDim1D, A.getTotalInDimSize()}}, - /*requireSurjective=*/false); - - SmallVector> retInDims; - SmallVector> retOutDims; - for (StringAttr dim : B.getInDimNames()) { - retInDims.push_back({dim, B.getInDimSize(dim)}); - } - for (StringAttr dim : A.getInDimNames()) { - retOutDims.push_back({dim, A.getInDimSize(dim)}); - } - return retFlattened.reshapeIns(retInDims).reshapeOuts(retOutDims); -} - -} // namespace - -LinearLayout LinearLayout::invertAndCompose(const LinearLayout &outer) const { - // TODO(Lezcano) Make friend and perhaps rename to `convertFrom` or `lstsq` - // For this, we need to implement our LLVM lowerings by inverting the "outer" - // layout, and then iterating over the elements from the "this" layout and - // fetching the corresponding element from the "outer" layout. This exercises - // the broadcasting that we incentivise via choosing the minimum norm solution - // in lstsq. - - // The order of dims does not matter. We choose to transpose outer - auto outDims = llvm::to_vector(getOutDimNames()); - assertDimsEqualIgnoringOrder(outDims, outer.getOutDimNames()); - const auto &B = *this; - const auto A = outer.transposeOuts(outDims); - for (auto dim : outDims) { - assert(A.getOutDimSize(dim) >= B.getOutDimSize(dim) && - ("A.invertAndCompose(B) called with incompatible output shapes in " + - dim.str() + ": " + std::to_string(A.getOutDimSize(dim)) + - " >= " + std::to_string(B.getOutDimSize(dim))) - .c_str()); - } - - // Broadcasting heuristic - // Imagine we have two layouts with `warps = [[0, 0],  [0, 0]]` - // (broadcasting) on both layouts. We could map any warp to any warp in the - // conversion. Now, we want to map them as the identity map, to mark that - // nothing needs to be done there (`lstsq` would map all the warps to the - // zero warp, minimum norm solution). The heuristic here is as follows: - // - If a dimension is the same for both layouts, we want to map it as the - // identity - // Equivalently, we don't add it to the conversion - // - Otherwise, we just call lstsq (i.e. map all the equivalent elements - // to the same input element) to take advantage of broadcasting in shared - // memory and avoid saving repeated elements in shared memory - - // FIXME: We should check that the other dimensions don't touch the image of - // this dimension. - SmallVector identityDims; - for (auto dim : A.getInDimNames()) { - if (B.hasInDim(dim) && - A.sublayout(dim, outDims) == B.sublayout(dim, outDims)) { - identityDims.push_back(dim); - } - } - SmallVector ANonIdentityInDims; - SmallVector BNonIdentityInDims; - for (auto dim : A.getInDimNames()) { - if (!llvm::is_contained(identityDims, dim)) { - ANonIdentityInDims.push_back(dim); - } - } - for (auto dim : B.getInDimNames()) { - if (!llvm::is_contained(identityDims, dim)) { - BNonIdentityInDims.push_back(dim); - } - } - - auto AReduced = A.sublayout(ANonIdentityInDims, outDims); - auto BReduced = B.sublayout(BNonIdentityInDims, outDims); - - // If one is empty, the other must be empty as well - assert((ANonIdentityInDims.empty()) == (BNonIdentityInDims.empty())); - bool isEmpty = ANonIdentityInDims.empty(); - - auto ret = isEmpty ? LinearLayout::empty() : lstsq(AReduced, BReduced); - - // TODO(Lezcano): We should return the reduced layout instead of re-adding the - // identity maps. With this, we'll be able to kill `minimalCvtLayout` - - // Add the identity maps for the dimensions that are the same for both layouts - for (auto dim : identityDims) { - ret *= LinearLayout::identity1D(A.getInDimSize(dim), dim, dim); - } - - // Reorder the dimensions in the result to match the order expected by the - // current and outer layouts. - return ret.transposeIns(llvm::to_vector(B.getInDimNames())) - .transposeOuts(llvm::to_vector(A.getInDimNames())); -} - -LinearLayout LinearLayout::invert() const { - assert(isInvertible() && - "A linear layout must be surjective and square to be invertible"); - return pseudoinvert(); -} - -LinearLayout LinearLayout::pseudoinvert() const { - LinearLayout identity = LinearLayout::empty(); - for (auto outDim : getOutDimNames()) { - identity *= LinearLayout::identity1D(getOutDimSize(outDim), outDim, outDim); - } - return identity.invertAndCompose(*this); -} - -LinearLayout LinearLayout::unsqueezeIn(StringAttr dim) const { - assert(getInDimSize(dim) == 1); - SmallVector> newInDims; - for (auto inDim : getInDimNames()) { - if (inDim != dim) { - newInDims.push_back({inDim, getInDimSize(inDim)}); - } - } - return reshapeIns(newInDims); -} - -LinearLayout LinearLayout::unsqueezeOut(StringAttr dim) const { - assert(getOutDimSize(dim) == 1); - SmallVector> newOutDims; - for (auto [outDim, outDimSize] : getOutDims()) { - if (outDim != dim) { - newOutDims.push_back({outDim, outDimSize}); - } - } - return LinearLayout(bases, newOutDims, isSurjective()); -} - -llvm::MapVector -LinearLayout::getFreeVariableMasks() const { - std::unique_ptr mat = getMatrix(*this); - int numRows = getTotalOutDimSizeLog2(); - int numCols = getTotalInDimSizeLog2(); - - // stride is specified in number of 64-bit words per row, and we pack our - // matrix so that there's only one uint64_t per row. - assert(numCols <= 64); - f2reduce::inplace_rref_strided(mat.get(), numRows, numCols, /*stride=*/1); - - // For each row in the RREF matrix, identify the column with the first "1". - // These columns correspond to the basic (i.e. non-free) variables. - std::set basicVars; - for (int r = 0; r < numRows; r++) { - if (mat[r] == 0) { - continue; - } - basicVars.insert(__builtin_ctzll(mat[r])); - } - - llvm::MapVector ret; - int c = 0; - for (StringAttr dim : getInDimNames()) { - int32_t mask = 0; - for (int i = 0; i < getInDimSizeLog2(dim); i++, c++) { - if (basicVars.count(c) == 0) { - mask |= (1 << i); - } - } - ret[dim] = mask; - } - return ret; -} - -LinearLayout LinearLayout::removeZeroBasesAlongDim(StringAttr stripDim) const { - LinearLayout::BasesT result; - for (auto &[inDim, inDimBases] : getBases()) { - auto &newInDimBases = result[inDim]; - if (inDim != stripDim) { - newInDimBases = inDimBases; - continue; - } - for (auto &basis : inDimBases) { - if (llvm::any_of(basis, [](int32_t val) { return val != 0; })) { - newInDimBases.push_back(basis); - } - } - } - SmallVector> newOutDimSizes; - for (auto outDim : getOutDimNames()) { - newOutDimSizes.push_back({outDim, getOutDimSize(outDim)}); - } - auto newLayout = LinearLayout(std::move(result), ArrayRef(newOutDimSizes), - this->isSurjective()); - return newLayout; -} - -size_t hash_value(const LinearLayout &layout) { - size_t seed = 0; - - // Hash the bases - for (const auto &base : layout.getBases()) { - // Hash the input dimension name - seed = llvm::hash_combine(seed, base.first); - - // Hash the vectors in bases - for (const auto &vec : base.second) { - for (int32_t val : vec) { - seed = llvm::hash_combine(seed, val); - } - } - } - - // Hash the output dimensions and their sizes - for (const auto &outDim : layout.getOutDimNames()) { - seed = llvm::hash_combine(seed, outDim, layout.getOutDimSize(outDim)); - } - // Don't hash the surjective flag as it's a cached property - return seed; -} - -bool operator==(const LinearLayout &lhs, const LinearLayout &rhs) { - if (!lhs.equalIgnoringOutDimSizes(rhs)) - return false; - - for (const auto &[lhsOutDimAndSize, rhsOutDimAndSize] : - llvm::zip(lhs.outDims, rhs.outDims)) { - if (lhsOutDimAndSize.second != rhsOutDimAndSize.second) - return false; - } - return true; -} - -bool LinearLayout::equalIgnoringOutDimSizes(const LinearLayout &other) const { - // llvm::MapVector doesn't have an operator== :(. - if (llvm::to_vector(this->getOutDimNames()) != - llvm::to_vector(other.getOutDimNames())) - return false; - if (this->bases.size() != other.bases.size()) - return false; - for (auto it1 = this->bases.begin(), it2 = other.bases.begin(); - it1 != this->bases.end(); ++it1, ++it2) { - if (*it1 != *it2) - return false; - } - return true; -} - -std::string LinearLayout::toString() const { - // Start with a newline because we print out a bulleted list; it doesn't - // make sense for the first line of this list to be on the same line as - // any previous text. - std::string ret = "\n"; - std::string outDimsStr = - "[" + - join(outDims, ", ", - [](auto dimAndSize) { - auto [outDim, size] = dimAndSize; - return outDim.str() + " (size " + std::to_string(size) + ")"; - }) + - "]"; - - if (bases.empty()) { - if (outDims.empty()) { - return "\n(empty layout)"; - } else { - return "\n(empty layout with out-dims " + outDimsStr + ")"; - } - } - - // TODO: Add spaces for alignment. - for (const auto &[inDim, inDimBases] : bases) { - if (inDimBases.empty()) { - ret += " - " + inDim.str() + " is a size 1 dimension\n"; - continue; - } - - ret += " - " + - join(llvm::seq(inDimBases.size()), "\n ", - [&, &inDim = inDim, &inDimBases = inDimBases](int i) { - return inDim.str() + "=" + std::to_string(1 << i) + " -> (" + - join(inDimBases[i], ", ") + ")"; - }) + - "\n"; - } - ret += "where out dims are: " + outDimsStr; - return ret; -} - -LinearLayout ColumnAction::apply(const LinearLayout &layout) const { - assert(layout.hasInDim(inDim)); - assert(layout.getInDimSizeLog2(inDim) == inSizeLog2 && - "Layout has a different size than the ColumnAction"); - if (m_isIdentity) { - return layout; - } - - auto bases = layout.getBases(); - const auto &basesInDim = bases[inDim]; - std::vector> newBases; - newBases.reserve(action.size()); - for (size_t a : action) { - newBases.push_back(basesInDim[a]); - } - bases[inDim] = std::move(newBases); - - SmallVector> outDims; - for (auto outDim : layout.getOutDimNames()) { - outDims.emplace_back(outDim, layout.getOutDimSize(outDim)); - } - return LinearLayout(std::move(bases), std::move(outDims), - /*requireSurjective=*/false); -} - -SmallVector ColumnAction::apply(ValueRange values) const { - assert(values.size() == (1 << inSizeLog2) && - "Values have a different size than the ColumnAction"); - assert(inDim.str() == "register" && "Values are in registers, so we can only " - "apply ColumnAction to registers"); - if (m_isIdentity) { - return values; - } - auto permLL = apply(LinearLayout::identity1D(values.size(), inDim, inDim)); - SmallVector ret; - ret.reserve(permLL.getInDimSize(inDim)); - for (int i = 0; i < permLL.getInDimSize(inDim); i++) { - int32_t srcIdx = permLL.apply({{inDim, i}}).begin()->second; - ret.push_back(values[srcIdx]); - } - return ret; -} - -ColumnAction ColumnAction::leftCompose(const ColumnAction &other) const { - assert(inDim == other.inDim); - assert(inSizeLog2 == other.inSizeLog2); - assert(action.size() == other.action.size()); - auto newAction = SmallVector(action.size()); - for (size_t i = 0; i < action.size(); i++) { - newAction[i] = action[other.action[i]]; - } - return ColumnAction(newAction, inDim, inSizeLog2); -} - -ColumnAction ColumnAction::inverse() const { - auto invPerm = SmallVector(action.size()); - for (size_t i = 0; i < action.size(); i++) { - invPerm[action[i]] = i; - } - return ColumnAction(invPerm, inDim, inSizeLog2); -} - -std::string ColumnAction::toString() const { - std::string ret = "ColumnAction(["; - ret += join(action, ", "); - ret += "], " + inDim.str() + ", " + std::to_string(inSizeLog2) + ")"; - return ret; -} - -// Build a matrix of size sum(outDimSizeLog2) x sum(inDimSizeLog2) representing -// the bases of the given layout. This can then be used by f2reduce. -// -// This function is called from the constructor of LinearLayout, so be careful -// not to use any functions that create LLs in here. -std::unique_ptr getMatrix(const LinearLayout &layout) { - int numRows = layout.getTotalOutDimSizeLog2(); - int numCols = layout.getTotalInDimSizeLog2(); - - // Don't handle giant LLs. This makes some things easier; for example, each - // row can be a single uint64_t. - assert(numCols <= 64 && "LinearLayout too large"); - assert(numRows <= 64 && "LinearLayout too large"); - - // Suppose we have a layout specified by the following values. - // - // L(0,1) = (0b01, 0b1) - // L(0,2) = (0b10, 0b0) - // L(1,0) = (0b10, 0b0) - // L(2,0) = (0b11, 0b0) - // - // We will create one column per entry above. The max bit width of the - // codomain is (2,1), so our matrix will have 2+1=3 rows. The final matrix - // will be - // - // | L(0,1)[0] L(0,2)[0] L(1,0)[0] L(2,0)[0] | | 0b1001 | - // | ↓ ↓ ↓ ↓ | | 0b0111 | - // | L(0,1)[1] L(0,2)[1] L(1,0)[1] L(2,0)[1] | = | 0b1000 | - // | ↓ ↓ ↓ ↓ | - // - // Note `new uint64_t[n]()` is zero-initialized, but `new uint64_t[n]` is not. - std::unique_ptr m(new uint64_t[numRows]()); - int r = 0; - for (StringAttr outDim : layout.getOutDimNames()) { - int c = 0; - for (StringAttr inDim : layout.getInDimNames()) { - for (int i = 0; i < layout.getInDimSizeLog2(inDim); i++) { - uint64_t basis = layout.getBasis(inDim, i, outDim); - for (int j = 0; j < layout.getOutDimSizeLog2(outDim); j++) { - m[r + j] |= ((basis >> j) & 1) << c; - } - c++; - } - } - r += layout.getOutDimSizeLog2(outDim); - } - - return m; -} - -} // namespace mlir::triton diff --git a/third_party/mthreads/include/triton/Analysis/Alias.h b/third_party/mthreads/spec_cpp/include/triton/Analysis/Alias.h similarity index 100% rename from third_party/mthreads/include/triton/Analysis/Alias.h rename to third_party/mthreads/spec_cpp/include/triton/Analysis/Alias.h diff --git a/third_party/mthreads/include/triton/Analysis/Allocation.h b/third_party/mthreads/spec_cpp/include/triton/Analysis/Allocation.h similarity index 100% rename from third_party/mthreads/include/triton/Analysis/Allocation.h rename to third_party/mthreads/spec_cpp/include/triton/Analysis/Allocation.h diff --git a/third_party/mthreads/include/triton/Analysis/BufferRegion.h b/third_party/mthreads/spec_cpp/include/triton/Analysis/BufferRegion.h similarity index 100% rename from third_party/mthreads/include/triton/Analysis/BufferRegion.h rename to third_party/mthreads/spec_cpp/include/triton/Analysis/BufferRegion.h diff --git a/third_party/mthreads/include/triton/Analysis/Membar.h b/third_party/mthreads/spec_cpp/include/triton/Analysis/Membar.h similarity index 100% rename from third_party/mthreads/include/triton/Analysis/Membar.h rename to third_party/mthreads/spec_cpp/include/triton/Analysis/Membar.h diff --git a/third_party/mthreads/include/triton/Analysis/Utility.h b/third_party/mthreads/spec_cpp/include/triton/Analysis/Utility.h similarity index 100% rename from third_party/mthreads/include/triton/Analysis/Utility.h rename to third_party/mthreads/spec_cpp/include/triton/Analysis/Utility.h diff --git a/third_party/mthreads/include/triton/Conversion/TritonGPUToLLVM/ElementwiseOpToLLVMBase.h b/third_party/mthreads/spec_cpp/include/triton/Conversion/TritonGPUToLLVM/ElementwiseOpToLLVMBase.h similarity index 100% rename from third_party/mthreads/include/triton/Conversion/TritonGPUToLLVM/ElementwiseOpToLLVMBase.h rename to third_party/mthreads/spec_cpp/include/triton/Conversion/TritonGPUToLLVM/ElementwiseOpToLLVMBase.h diff --git a/third_party/mthreads/include/triton/Conversion/TritonGPUToLLVM/PatternTritonGPUOpToLLVM.h b/third_party/mthreads/spec_cpp/include/triton/Conversion/TritonGPUToLLVM/PatternTritonGPUOpToLLVM.h similarity index 100% rename from third_party/mthreads/include/triton/Conversion/TritonGPUToLLVM/PatternTritonGPUOpToLLVM.h rename to third_party/mthreads/spec_cpp/include/triton/Conversion/TritonGPUToLLVM/PatternTritonGPUOpToLLVM.h diff --git a/third_party/mthreads/include/triton/Conversion/TritonGPUToLLVM/TargetInfoBase.h b/third_party/mthreads/spec_cpp/include/triton/Conversion/TritonGPUToLLVM/TargetInfoBase.h similarity index 100% rename from third_party/mthreads/include/triton/Conversion/TritonGPUToLLVM/TargetInfoBase.h rename to third_party/mthreads/spec_cpp/include/triton/Conversion/TritonGPUToLLVM/TargetInfoBase.h diff --git a/third_party/mthreads/include/triton/Conversion/TritonGPUToLLVM/Utility.h b/third_party/mthreads/spec_cpp/include/triton/Conversion/TritonGPUToLLVM/Utility.h similarity index 100% rename from third_party/mthreads/include/triton/Conversion/TritonGPUToLLVM/Utility.h rename to third_party/mthreads/spec_cpp/include/triton/Conversion/TritonGPUToLLVM/Utility.h diff --git a/third_party/mthreads/include/triton/Conversion/TritonGPUToLLVM/WarpSpecializeUtility.h b/third_party/mthreads/spec_cpp/include/triton/Conversion/TritonGPUToLLVM/WarpSpecializeUtility.h similarity index 100% rename from third_party/mthreads/include/triton/Conversion/TritonGPUToLLVM/WarpSpecializeUtility.h rename to third_party/mthreads/spec_cpp/include/triton/Conversion/TritonGPUToLLVM/WarpSpecializeUtility.h diff --git a/third_party/mthreads/include/triton/Dialect/Gluon/Transforms/Passes.td b/third_party/mthreads/spec_cpp/include/triton/Dialect/Gluon/Transforms/Passes.td similarity index 100% rename from third_party/mthreads/include/triton/Dialect/Gluon/Transforms/Passes.td rename to third_party/mthreads/spec_cpp/include/triton/Dialect/Gluon/Transforms/Passes.td diff --git a/third_party/mthreads/include/triton/Dialect/NVGPU/IR/CMakeLists.txt b/third_party/mthreads/spec_cpp/include/triton/Dialect/NVGPU/IR/CMakeLists.txt similarity index 100% rename from third_party/mthreads/include/triton/Dialect/NVGPU/IR/CMakeLists.txt rename to third_party/mthreads/spec_cpp/include/triton/Dialect/NVGPU/IR/CMakeLists.txt diff --git a/third_party/mthreads/include/triton/Dialect/NVGPU/IR/Dialect.h b/third_party/mthreads/spec_cpp/include/triton/Dialect/NVGPU/IR/Dialect.h similarity index 100% rename from third_party/mthreads/include/triton/Dialect/NVGPU/IR/Dialect.h rename to third_party/mthreads/spec_cpp/include/triton/Dialect/NVGPU/IR/Dialect.h diff --git a/third_party/mthreads/include/triton/Dialect/NVGPU/IR/NVGPUAttrDefs.td b/third_party/mthreads/spec_cpp/include/triton/Dialect/NVGPU/IR/NVGPUAttrDefs.td similarity index 100% rename from third_party/mthreads/include/triton/Dialect/NVGPU/IR/NVGPUAttrDefs.td rename to third_party/mthreads/spec_cpp/include/triton/Dialect/NVGPU/IR/NVGPUAttrDefs.td diff --git a/third_party/mthreads/include/triton/Dialect/NVGPU/IR/NVGPUDialect.td b/third_party/mthreads/spec_cpp/include/triton/Dialect/NVGPU/IR/NVGPUDialect.td similarity index 100% rename from third_party/mthreads/include/triton/Dialect/NVGPU/IR/NVGPUDialect.td rename to third_party/mthreads/spec_cpp/include/triton/Dialect/NVGPU/IR/NVGPUDialect.td diff --git a/third_party/mthreads/include/triton/Dialect/NVGPU/IR/NVGPUOps.td b/third_party/mthreads/spec_cpp/include/triton/Dialect/NVGPU/IR/NVGPUOps.td similarity index 100% rename from third_party/mthreads/include/triton/Dialect/NVGPU/IR/NVGPUOps.td rename to third_party/mthreads/spec_cpp/include/triton/Dialect/NVGPU/IR/NVGPUOps.td diff --git a/third_party/mthreads/include/triton/Dialect/NVWS/IR/CMakeLists.txt b/third_party/mthreads/spec_cpp/include/triton/Dialect/NVWS/IR/CMakeLists.txt similarity index 100% rename from third_party/mthreads/include/triton/Dialect/NVWS/IR/CMakeLists.txt rename to third_party/mthreads/spec_cpp/include/triton/Dialect/NVWS/IR/CMakeLists.txt diff --git a/third_party/mthreads/include/triton/Dialect/NVWS/IR/Dialect.h b/third_party/mthreads/spec_cpp/include/triton/Dialect/NVWS/IR/Dialect.h similarity index 100% rename from third_party/mthreads/include/triton/Dialect/NVWS/IR/Dialect.h rename to third_party/mthreads/spec_cpp/include/triton/Dialect/NVWS/IR/Dialect.h diff --git a/third_party/mthreads/include/triton/Dialect/NVWS/IR/NVWSAttrDefs.td b/third_party/mthreads/spec_cpp/include/triton/Dialect/NVWS/IR/NVWSAttrDefs.td similarity index 100% rename from third_party/mthreads/include/triton/Dialect/NVWS/IR/NVWSAttrDefs.td rename to third_party/mthreads/spec_cpp/include/triton/Dialect/NVWS/IR/NVWSAttrDefs.td diff --git a/third_party/mthreads/include/triton/Dialect/NVWS/IR/NVWSDialect.td b/third_party/mthreads/spec_cpp/include/triton/Dialect/NVWS/IR/NVWSDialect.td similarity index 100% rename from third_party/mthreads/include/triton/Dialect/NVWS/IR/NVWSDialect.td rename to third_party/mthreads/spec_cpp/include/triton/Dialect/NVWS/IR/NVWSDialect.td diff --git a/third_party/mthreads/include/triton/Dialect/NVWS/IR/NVWSOpInterfaces.td b/third_party/mthreads/spec_cpp/include/triton/Dialect/NVWS/IR/NVWSOpInterfaces.td similarity index 100% rename from third_party/mthreads/include/triton/Dialect/NVWS/IR/NVWSOpInterfaces.td rename to third_party/mthreads/spec_cpp/include/triton/Dialect/NVWS/IR/NVWSOpInterfaces.td diff --git a/third_party/mthreads/include/triton/Dialect/NVWS/IR/NVWSOps.td b/third_party/mthreads/spec_cpp/include/triton/Dialect/NVWS/IR/NVWSOps.td similarity index 100% rename from third_party/mthreads/include/triton/Dialect/NVWS/IR/NVWSOps.td rename to third_party/mthreads/spec_cpp/include/triton/Dialect/NVWS/IR/NVWSOps.td diff --git a/third_party/mthreads/include/triton/Dialect/NVWS/IR/NVWSTypes.td b/third_party/mthreads/spec_cpp/include/triton/Dialect/NVWS/IR/NVWSTypes.td similarity index 100% rename from third_party/mthreads/include/triton/Dialect/NVWS/IR/NVWSTypes.td rename to third_party/mthreads/spec_cpp/include/triton/Dialect/NVWS/IR/NVWSTypes.td diff --git a/third_party/mthreads/include/triton/Dialect/NVWS/Transforms/CMakeLists.txt b/third_party/mthreads/spec_cpp/include/triton/Dialect/NVWS/Transforms/CMakeLists.txt similarity index 100% rename from third_party/mthreads/include/triton/Dialect/NVWS/Transforms/CMakeLists.txt rename to third_party/mthreads/spec_cpp/include/triton/Dialect/NVWS/Transforms/CMakeLists.txt diff --git a/third_party/mthreads/include/triton/Dialect/NVWS/Transforms/Passes.h b/third_party/mthreads/spec_cpp/include/triton/Dialect/NVWS/Transforms/Passes.h similarity index 100% rename from third_party/mthreads/include/triton/Dialect/NVWS/Transforms/Passes.h rename to third_party/mthreads/spec_cpp/include/triton/Dialect/NVWS/Transforms/Passes.h diff --git a/third_party/mthreads/include/triton/Dialect/NVWS/Transforms/Passes.td b/third_party/mthreads/spec_cpp/include/triton/Dialect/NVWS/Transforms/Passes.td similarity index 100% rename from third_party/mthreads/include/triton/Dialect/NVWS/Transforms/Passes.td rename to third_party/mthreads/spec_cpp/include/triton/Dialect/NVWS/Transforms/Passes.td diff --git a/third_party/mthreads/spec_cpp/include/triton/Dialect/Triton/IR/CMakeLists.txt b/third_party/mthreads/spec_cpp/include/triton/Dialect/Triton/IR/CMakeLists.txt new file mode 100644 index 0000000000..c84bcba83f --- /dev/null +++ b/third_party/mthreads/spec_cpp/include/triton/Dialect/Triton/IR/CMakeLists.txt @@ -0,0 +1,4 @@ +set(LLVM_TARGET_DEFINITIONS TritonTypeInterfaces.td) +mlir_tablegen(TypeInterfaces.h.inc -gen-type-interface-decls) +mlir_tablegen(TypeInterfaces.cpp.inc -gen-type-interface-defs) +add_public_tablegen_target(TritonTypeInterfacesIncGen) diff --git a/third_party/mthreads/include/triton/Dialect/Triton/IR/Dialect.h b/third_party/mthreads/spec_cpp/include/triton/Dialect/Triton/IR/Dialect.h similarity index 100% rename from third_party/mthreads/include/triton/Dialect/Triton/IR/Dialect.h rename to third_party/mthreads/spec_cpp/include/triton/Dialect/Triton/IR/Dialect.h diff --git a/third_party/mthreads/include/triton/Dialect/Triton/IR/Traits.h b/third_party/mthreads/spec_cpp/include/triton/Dialect/Triton/IR/Traits.h similarity index 100% rename from third_party/mthreads/include/triton/Dialect/Triton/IR/Traits.h rename to third_party/mthreads/spec_cpp/include/triton/Dialect/Triton/IR/Traits.h diff --git a/third_party/mthreads/include/triton/Dialect/Triton/IR/TritonInterfaces.td b/third_party/mthreads/spec_cpp/include/triton/Dialect/Triton/IR/TritonInterfaces.td similarity index 100% rename from third_party/mthreads/include/triton/Dialect/Triton/IR/TritonInterfaces.td rename to third_party/mthreads/spec_cpp/include/triton/Dialect/Triton/IR/TritonInterfaces.td diff --git a/third_party/mthreads/include/triton/Dialect/Triton/IR/TritonOps.td b/third_party/mthreads/spec_cpp/include/triton/Dialect/Triton/IR/TritonOps.td similarity index 100% rename from third_party/mthreads/include/triton/Dialect/Triton/IR/TritonOps.td rename to third_party/mthreads/spec_cpp/include/triton/Dialect/Triton/IR/TritonOps.td diff --git a/third_party/mthreads/include/triton/Dialect/Triton/IR/TritonTypeInterfaces.td b/third_party/mthreads/spec_cpp/include/triton/Dialect/Triton/IR/TritonTypeInterfaces.td similarity index 100% rename from third_party/mthreads/include/triton/Dialect/Triton/IR/TritonTypeInterfaces.td rename to third_party/mthreads/spec_cpp/include/triton/Dialect/Triton/IR/TritonTypeInterfaces.td diff --git a/third_party/mthreads/include/triton/Dialect/Triton/IR/TritonTypes.td b/third_party/mthreads/spec_cpp/include/triton/Dialect/Triton/IR/TritonTypes.td similarity index 100% rename from third_party/mthreads/include/triton/Dialect/Triton/IR/TritonTypes.td rename to third_party/mthreads/spec_cpp/include/triton/Dialect/Triton/IR/TritonTypes.td diff --git a/third_party/mthreads/include/triton/Dialect/Triton/IR/Types.h b/third_party/mthreads/spec_cpp/include/triton/Dialect/Triton/IR/Types.h similarity index 100% rename from third_party/mthreads/include/triton/Dialect/Triton/IR/Types.h rename to third_party/mthreads/spec_cpp/include/triton/Dialect/Triton/IR/Types.h diff --git a/third_party/mthreads/include/triton/Dialect/Triton/IR/Utility.h b/third_party/mthreads/spec_cpp/include/triton/Dialect/Triton/IR/Utility.h similarity index 100% rename from third_party/mthreads/include/triton/Dialect/Triton/IR/Utility.h rename to third_party/mthreads/spec_cpp/include/triton/Dialect/Triton/IR/Utility.h diff --git a/third_party/mthreads/include/triton/Dialect/TritonGPU/IR/Attributes.h b/third_party/mthreads/spec_cpp/include/triton/Dialect/TritonGPU/IR/Attributes.h similarity index 100% rename from third_party/mthreads/include/triton/Dialect/TritonGPU/IR/Attributes.h rename to third_party/mthreads/spec_cpp/include/triton/Dialect/TritonGPU/IR/Attributes.h diff --git a/third_party/mthreads/include/triton/Dialect/TritonGPU/IR/CGAEncodingAttr.h b/third_party/mthreads/spec_cpp/include/triton/Dialect/TritonGPU/IR/CGAEncodingAttr.h similarity index 100% rename from third_party/mthreads/include/triton/Dialect/TritonGPU/IR/CGAEncodingAttr.h rename to third_party/mthreads/spec_cpp/include/triton/Dialect/TritonGPU/IR/CGAEncodingAttr.h diff --git a/third_party/mthreads/include/triton/Dialect/TritonGPU/IR/CGAEncodingAttr.td b/third_party/mthreads/spec_cpp/include/triton/Dialect/TritonGPU/IR/CGAEncodingAttr.td similarity index 100% rename from third_party/mthreads/include/triton/Dialect/TritonGPU/IR/CGAEncodingAttr.td rename to third_party/mthreads/spec_cpp/include/triton/Dialect/TritonGPU/IR/CGAEncodingAttr.td diff --git a/third_party/mthreads/spec_cpp/include/triton/Dialect/TritonGPU/IR/CMakeLists.txt b/third_party/mthreads/spec_cpp/include/triton/Dialect/TritonGPU/IR/CMakeLists.txt new file mode 100644 index 0000000000..593afbf5e3 --- /dev/null +++ b/third_party/mthreads/spec_cpp/include/triton/Dialect/TritonGPU/IR/CMakeLists.txt @@ -0,0 +1,8 @@ +set(LLVM_TARGET_DEFINITIONS CGAEncodingAttr.td) +mlir_tablegen(CGAEncodingAttr.h.inc -gen-attrdef-decls) +add_public_tablegen_target(TritonGPUCGAAttrIncGen) + +set(LLVM_TARGET_DEFINITIONS TritonGPUEnums.td) +mlir_tablegen(OpsEnums.h.inc -gen-enum-decls) +mlir_tablegen(OpsEnums.cpp.inc -gen-enum-defs) +add_public_tablegen_target(TritonGPUOpsEnumsIncGen) diff --git a/third_party/mthreads/include/triton/Dialect/TritonGPU/IR/Dialect.h b/third_party/mthreads/spec_cpp/include/triton/Dialect/TritonGPU/IR/Dialect.h similarity index 100% rename from third_party/mthreads/include/triton/Dialect/TritonGPU/IR/Dialect.h rename to third_party/mthreads/spec_cpp/include/triton/Dialect/TritonGPU/IR/Dialect.h diff --git a/third_party/mthreads/include/triton/Dialect/TritonGPU/IR/LinearLayoutConversions.h b/third_party/mthreads/spec_cpp/include/triton/Dialect/TritonGPU/IR/LinearLayoutConversions.h similarity index 100% rename from third_party/mthreads/include/triton/Dialect/TritonGPU/IR/LinearLayoutConversions.h rename to third_party/mthreads/spec_cpp/include/triton/Dialect/TritonGPU/IR/LinearLayoutConversions.h diff --git a/third_party/mthreads/include/triton/Dialect/TritonGPU/IR/Traits.h b/third_party/mthreads/spec_cpp/include/triton/Dialect/TritonGPU/IR/Traits.h similarity index 100% rename from third_party/mthreads/include/triton/Dialect/TritonGPU/IR/Traits.h rename to third_party/mthreads/spec_cpp/include/triton/Dialect/TritonGPU/IR/Traits.h diff --git a/third_party/mthreads/include/triton/Dialect/TritonGPU/IR/TritonGPUAttrBase.td b/third_party/mthreads/spec_cpp/include/triton/Dialect/TritonGPU/IR/TritonGPUAttrBase.td similarity index 100% rename from third_party/mthreads/include/triton/Dialect/TritonGPU/IR/TritonGPUAttrBase.td rename to third_party/mthreads/spec_cpp/include/triton/Dialect/TritonGPU/IR/TritonGPUAttrBase.td diff --git a/third_party/mthreads/include/triton/Dialect/TritonGPU/IR/TritonGPUAttrDefs.td b/third_party/mthreads/spec_cpp/include/triton/Dialect/TritonGPU/IR/TritonGPUAttrDefs.td similarity index 100% rename from third_party/mthreads/include/triton/Dialect/TritonGPU/IR/TritonGPUAttrDefs.td rename to third_party/mthreads/spec_cpp/include/triton/Dialect/TritonGPU/IR/TritonGPUAttrDefs.td diff --git a/third_party/mthreads/include/triton/Dialect/TritonGPU/IR/TritonGPUAttrImpls.td b/third_party/mthreads/spec_cpp/include/triton/Dialect/TritonGPU/IR/TritonGPUAttrImpls.td similarity index 100% rename from third_party/mthreads/include/triton/Dialect/TritonGPU/IR/TritonGPUAttrImpls.td rename to third_party/mthreads/spec_cpp/include/triton/Dialect/TritonGPU/IR/TritonGPUAttrImpls.td diff --git a/third_party/mthreads/include/triton/Dialect/TritonGPU/IR/TritonGPUEnums.td b/third_party/mthreads/spec_cpp/include/triton/Dialect/TritonGPU/IR/TritonGPUEnums.td similarity index 100% rename from third_party/mthreads/include/triton/Dialect/TritonGPU/IR/TritonGPUEnums.td rename to third_party/mthreads/spec_cpp/include/triton/Dialect/TritonGPU/IR/TritonGPUEnums.td diff --git a/third_party/mthreads/include/triton/Dialect/TritonGPU/IR/TritonGPUInterfaces.h b/third_party/mthreads/spec_cpp/include/triton/Dialect/TritonGPU/IR/TritonGPUInterfaces.h similarity index 100% rename from third_party/mthreads/include/triton/Dialect/TritonGPU/IR/TritonGPUInterfaces.h rename to third_party/mthreads/spec_cpp/include/triton/Dialect/TritonGPU/IR/TritonGPUInterfaces.h diff --git a/third_party/mthreads/include/triton/Dialect/TritonGPU/IR/TritonGPUOps.td b/third_party/mthreads/spec_cpp/include/triton/Dialect/TritonGPU/IR/TritonGPUOps.td similarity index 100% rename from third_party/mthreads/include/triton/Dialect/TritonGPU/IR/TritonGPUOps.td rename to third_party/mthreads/spec_cpp/include/triton/Dialect/TritonGPU/IR/TritonGPUOps.td diff --git a/third_party/mthreads/include/triton/Dialect/TritonGPU/IR/TritonGPUTypes.td b/third_party/mthreads/spec_cpp/include/triton/Dialect/TritonGPU/IR/TritonGPUTypes.td similarity index 100% rename from third_party/mthreads/include/triton/Dialect/TritonGPU/IR/TritonGPUTypes.td rename to third_party/mthreads/spec_cpp/include/triton/Dialect/TritonGPU/IR/TritonGPUTypes.td diff --git a/third_party/mthreads/include/triton/Dialect/TritonGPU/Transforms/CoalesceUtils.h b/third_party/mthreads/spec_cpp/include/triton/Dialect/TritonGPU/Transforms/CoalesceUtils.h similarity index 100% rename from third_party/mthreads/include/triton/Dialect/TritonGPU/Transforms/CoalesceUtils.h rename to third_party/mthreads/spec_cpp/include/triton/Dialect/TritonGPU/Transforms/CoalesceUtils.h diff --git a/third_party/mthreads/include/triton/Dialect/TritonGPU/Transforms/PartitionSchedulingUtility.h b/third_party/mthreads/spec_cpp/include/triton/Dialect/TritonGPU/Transforms/PartitionSchedulingUtility.h similarity index 100% rename from third_party/mthreads/include/triton/Dialect/TritonGPU/Transforms/PartitionSchedulingUtility.h rename to third_party/mthreads/spec_cpp/include/triton/Dialect/TritonGPU/Transforms/PartitionSchedulingUtility.h diff --git a/third_party/mthreads/include/triton/Dialect/TritonGPU/Transforms/Passes.h b/third_party/mthreads/spec_cpp/include/triton/Dialect/TritonGPU/Transforms/Passes.h similarity index 100% rename from third_party/mthreads/include/triton/Dialect/TritonGPU/Transforms/Passes.h rename to third_party/mthreads/spec_cpp/include/triton/Dialect/TritonGPU/Transforms/Passes.h diff --git a/third_party/mthreads/include/triton/Dialect/TritonGPU/Transforms/Passes.td b/third_party/mthreads/spec_cpp/include/triton/Dialect/TritonGPU/Transforms/Passes.td similarity index 100% rename from third_party/mthreads/include/triton/Dialect/TritonGPU/Transforms/Passes.td rename to third_party/mthreads/spec_cpp/include/triton/Dialect/TritonGPU/Transforms/Passes.td diff --git a/third_party/mthreads/include/triton/Dialect/TritonGPU/Transforms/Schedule.h b/third_party/mthreads/spec_cpp/include/triton/Dialect/TritonGPU/Transforms/Schedule.h similarity index 100% rename from third_party/mthreads/include/triton/Dialect/TritonGPU/Transforms/Schedule.h rename to third_party/mthreads/spec_cpp/include/triton/Dialect/TritonGPU/Transforms/Schedule.h diff --git a/third_party/mthreads/include/triton/Dialect/TritonGPU/Transforms/TritonGPUConversion.h b/third_party/mthreads/spec_cpp/include/triton/Dialect/TritonGPU/Transforms/TritonGPUConversion.h similarity index 100% rename from third_party/mthreads/include/triton/Dialect/TritonGPU/Transforms/TritonGPUConversion.h rename to third_party/mthreads/spec_cpp/include/triton/Dialect/TritonGPU/Transforms/TritonGPUConversion.h diff --git a/third_party/mthreads/include/triton/Dialect/TritonGPU/Transforms/Utility.h b/third_party/mthreads/spec_cpp/include/triton/Dialect/TritonGPU/Transforms/Utility.h similarity index 100% rename from third_party/mthreads/include/triton/Dialect/TritonGPU/Transforms/Utility.h rename to third_party/mthreads/spec_cpp/include/triton/Dialect/TritonGPU/Transforms/Utility.h diff --git a/third_party/mthreads/include/triton/Dialect/TritonInstrument/IR/FunctionBuilder.h b/third_party/mthreads/spec_cpp/include/triton/Dialect/TritonInstrument/IR/FunctionBuilder.h similarity index 100% rename from third_party/mthreads/include/triton/Dialect/TritonInstrument/IR/FunctionBuilder.h rename to third_party/mthreads/spec_cpp/include/triton/Dialect/TritonInstrument/IR/FunctionBuilder.h diff --git a/third_party/mthreads/include/triton/Dialect/TritonInstrument/IR/TritonInstrumentOps.td b/third_party/mthreads/spec_cpp/include/triton/Dialect/TritonInstrument/IR/TritonInstrumentOps.td similarity index 100% rename from third_party/mthreads/include/triton/Dialect/TritonInstrument/IR/TritonInstrumentOps.td rename to third_party/mthreads/spec_cpp/include/triton/Dialect/TritonInstrument/IR/TritonInstrumentOps.td diff --git a/third_party/mthreads/include/triton/Dialect/TritonInstrument/IR/Utility.h b/third_party/mthreads/spec_cpp/include/triton/Dialect/TritonInstrument/IR/Utility.h similarity index 100% rename from third_party/mthreads/include/triton/Dialect/TritonInstrument/IR/Utility.h rename to third_party/mthreads/spec_cpp/include/triton/Dialect/TritonInstrument/IR/Utility.h diff --git a/third_party/mthreads/spec_cpp/include/triton/Dialect/TritonNvidiaGPU/IR/CMakeLists.txt b/third_party/mthreads/spec_cpp/include/triton/Dialect/TritonNvidiaGPU/IR/CMakeLists.txt new file mode 100644 index 0000000000..b3d6c6e989 --- /dev/null +++ b/third_party/mthreads/spec_cpp/include/triton/Dialect/TritonNvidiaGPU/IR/CMakeLists.txt @@ -0,0 +1,4 @@ +set(LLVM_TARGET_DEFINITIONS TritonNvidiaGPUTypes.td) +mlir_tablegen(Types.h.inc -gen-typedef-decls) +mlir_tablegen(Types.cpp.inc -gen-typedef-defs) +add_public_tablegen_target(TritonNvidiaGPUTypesIncGen) diff --git a/third_party/mthreads/include/triton/Dialect/TritonNvidiaGPU/IR/Dialect.h b/third_party/mthreads/spec_cpp/include/triton/Dialect/TritonNvidiaGPU/IR/Dialect.h similarity index 100% rename from third_party/mthreads/include/triton/Dialect/TritonNvidiaGPU/IR/Dialect.h rename to third_party/mthreads/spec_cpp/include/triton/Dialect/TritonNvidiaGPU/IR/Dialect.h diff --git a/third_party/mthreads/include/triton/Dialect/TritonNvidiaGPU/IR/TritonNvidiaGPUAttrDefs.td b/third_party/mthreads/spec_cpp/include/triton/Dialect/TritonNvidiaGPU/IR/TritonNvidiaGPUAttrDefs.td similarity index 100% rename from third_party/mthreads/include/triton/Dialect/TritonNvidiaGPU/IR/TritonNvidiaGPUAttrDefs.td rename to third_party/mthreads/spec_cpp/include/triton/Dialect/TritonNvidiaGPU/IR/TritonNvidiaGPUAttrDefs.td diff --git a/third_party/mthreads/include/triton/Dialect/TritonNvidiaGPU/IR/TritonNvidiaGPUDialect.td b/third_party/mthreads/spec_cpp/include/triton/Dialect/TritonNvidiaGPU/IR/TritonNvidiaGPUDialect.td similarity index 100% rename from third_party/mthreads/include/triton/Dialect/TritonNvidiaGPU/IR/TritonNvidiaGPUDialect.td rename to third_party/mthreads/spec_cpp/include/triton/Dialect/TritonNvidiaGPU/IR/TritonNvidiaGPUDialect.td diff --git a/third_party/mthreads/include/triton/Dialect/TritonNvidiaGPU/IR/TritonNvidiaGPUOpInterfaces.td b/third_party/mthreads/spec_cpp/include/triton/Dialect/TritonNvidiaGPU/IR/TritonNvidiaGPUOpInterfaces.td similarity index 100% rename from third_party/mthreads/include/triton/Dialect/TritonNvidiaGPU/IR/TritonNvidiaGPUOpInterfaces.td rename to third_party/mthreads/spec_cpp/include/triton/Dialect/TritonNvidiaGPU/IR/TritonNvidiaGPUOpInterfaces.td diff --git a/third_party/mthreads/include/triton/Dialect/TritonNvidiaGPU/IR/TritonNvidiaGPUOps.td b/third_party/mthreads/spec_cpp/include/triton/Dialect/TritonNvidiaGPU/IR/TritonNvidiaGPUOps.td similarity index 100% rename from third_party/mthreads/include/triton/Dialect/TritonNvidiaGPU/IR/TritonNvidiaGPUOps.td rename to third_party/mthreads/spec_cpp/include/triton/Dialect/TritonNvidiaGPU/IR/TritonNvidiaGPUOps.td diff --git a/third_party/mthreads/include/triton/Dialect/TritonNvidiaGPU/IR/TritonNvidiaGPUTypes.td b/third_party/mthreads/spec_cpp/include/triton/Dialect/TritonNvidiaGPU/IR/TritonNvidiaGPUTypes.td similarity index 100% rename from third_party/mthreads/include/triton/Dialect/TritonNvidiaGPU/IR/TritonNvidiaGPUTypes.td rename to third_party/mthreads/spec_cpp/include/triton/Dialect/TritonNvidiaGPU/IR/TritonNvidiaGPUTypes.td diff --git a/third_party/mthreads/include/triton/Dialect/TritonNvidiaGPU/Transforms/TMAUtilities.h b/third_party/mthreads/spec_cpp/include/triton/Dialect/TritonNvidiaGPU/Transforms/TMAUtilities.h similarity index 100% rename from third_party/mthreads/include/triton/Dialect/TritonNvidiaGPU/Transforms/TMAUtilities.h rename to third_party/mthreads/spec_cpp/include/triton/Dialect/TritonNvidiaGPU/Transforms/TMAUtilities.h diff --git a/third_party/mthreads/include/triton/Tools/LinearLayout.h b/third_party/mthreads/spec_cpp/include/triton/Tools/LinearLayout.h similarity index 100% rename from third_party/mthreads/include/triton/Tools/LinearLayout.h rename to third_party/mthreads/spec_cpp/include/triton/Tools/LinearLayout.h diff --git a/third_party/mthreads/include/triton/Tools/PluginUtils.h b/third_party/mthreads/spec_cpp/include/triton/Tools/PluginUtils.h similarity index 100% rename from third_party/mthreads/include/triton/Tools/PluginUtils.h rename to third_party/mthreads/spec_cpp/include/triton/Tools/PluginUtils.h diff --git a/third_party/mthreads/include/triton/Tools/Sys/GetEnv.hpp b/third_party/mthreads/spec_cpp/include/triton/Tools/Sys/GetEnv.hpp similarity index 100% rename from third_party/mthreads/include/triton/Tools/Sys/GetEnv.hpp rename to third_party/mthreads/spec_cpp/include/triton/Tools/Sys/GetEnv.hpp diff --git a/third_party/mthreads/lib/Analysis/Alias.cpp b/third_party/mthreads/spec_cpp/lib/Analysis/Alias.cpp similarity index 100% rename from third_party/mthreads/lib/Analysis/Alias.cpp rename to third_party/mthreads/spec_cpp/lib/Analysis/Alias.cpp diff --git a/third_party/mthreads/lib/Analysis/Allocation.cpp b/third_party/mthreads/spec_cpp/lib/Analysis/Allocation.cpp similarity index 100% rename from third_party/mthreads/lib/Analysis/Allocation.cpp rename to third_party/mthreads/spec_cpp/lib/Analysis/Allocation.cpp diff --git a/third_party/mthreads/lib/Analysis/AxisInfo.cpp b/third_party/mthreads/spec_cpp/lib/Analysis/AxisInfo.cpp similarity index 100% rename from third_party/mthreads/lib/Analysis/AxisInfo.cpp rename to third_party/mthreads/spec_cpp/lib/Analysis/AxisInfo.cpp diff --git a/third_party/mthreads/lib/Analysis/BufferRegion.cpp b/third_party/mthreads/spec_cpp/lib/Analysis/BufferRegion.cpp similarity index 100% rename from third_party/mthreads/lib/Analysis/BufferRegion.cpp rename to third_party/mthreads/spec_cpp/lib/Analysis/BufferRegion.cpp diff --git a/third_party/mthreads/lib/Analysis/Membar.cpp b/third_party/mthreads/spec_cpp/lib/Analysis/Membar.cpp similarity index 100% rename from third_party/mthreads/lib/Analysis/Membar.cpp rename to third_party/mthreads/spec_cpp/lib/Analysis/Membar.cpp diff --git a/third_party/mthreads/lib/Analysis/Utility.cpp b/third_party/mthreads/spec_cpp/lib/Analysis/Utility.cpp similarity index 100% rename from third_party/mthreads/lib/Analysis/Utility.cpp rename to third_party/mthreads/spec_cpp/lib/Analysis/Utility.cpp diff --git a/third_party/mthreads/lib/Conversion/TritonGPUToLLVM/AllocateWarpGroups.cpp b/third_party/mthreads/spec_cpp/lib/Conversion/TritonGPUToLLVM/AllocateWarpGroups.cpp similarity index 100% rename from third_party/mthreads/lib/Conversion/TritonGPUToLLVM/AllocateWarpGroups.cpp rename to third_party/mthreads/spec_cpp/lib/Conversion/TritonGPUToLLVM/AllocateWarpGroups.cpp diff --git a/third_party/mthreads/lib/Conversion/TritonGPUToLLVM/AssertOpToLLVM.cpp b/third_party/mthreads/spec_cpp/lib/Conversion/TritonGPUToLLVM/AssertOpToLLVM.cpp similarity index 100% rename from third_party/mthreads/lib/Conversion/TritonGPUToLLVM/AssertOpToLLVM.cpp rename to third_party/mthreads/spec_cpp/lib/Conversion/TritonGPUToLLVM/AssertOpToLLVM.cpp diff --git a/third_party/mthreads/lib/Conversion/TritonGPUToLLVM/ConvertLayoutOpToLLVM.cpp b/third_party/mthreads/spec_cpp/lib/Conversion/TritonGPUToLLVM/ConvertLayoutOpToLLVM.cpp similarity index 100% rename from third_party/mthreads/lib/Conversion/TritonGPUToLLVM/ConvertLayoutOpToLLVM.cpp rename to third_party/mthreads/spec_cpp/lib/Conversion/TritonGPUToLLVM/ConvertLayoutOpToLLVM.cpp diff --git a/third_party/mthreads/lib/Conversion/TritonGPUToLLVM/DotOpToLLVM/FMA.cpp b/third_party/mthreads/spec_cpp/lib/Conversion/TritonGPUToLLVM/DotOpToLLVM/FMA.cpp similarity index 100% rename from third_party/mthreads/lib/Conversion/TritonGPUToLLVM/DotOpToLLVM/FMA.cpp rename to third_party/mthreads/spec_cpp/lib/Conversion/TritonGPUToLLVM/DotOpToLLVM/FMA.cpp diff --git a/third_party/mthreads/lib/Conversion/TritonGPUToLLVM/ElementwiseOpToLLVM.cpp b/third_party/mthreads/spec_cpp/lib/Conversion/TritonGPUToLLVM/ElementwiseOpToLLVM.cpp similarity index 100% rename from third_party/mthreads/lib/Conversion/TritonGPUToLLVM/ElementwiseOpToLLVM.cpp rename to third_party/mthreads/spec_cpp/lib/Conversion/TritonGPUToLLVM/ElementwiseOpToLLVM.cpp diff --git a/third_party/mthreads/lib/Conversion/TritonGPUToLLVM/FuncOpToLLVM.cpp b/third_party/mthreads/spec_cpp/lib/Conversion/TritonGPUToLLVM/FuncOpToLLVM.cpp similarity index 100% rename from third_party/mthreads/lib/Conversion/TritonGPUToLLVM/FuncOpToLLVM.cpp rename to third_party/mthreads/spec_cpp/lib/Conversion/TritonGPUToLLVM/FuncOpToLLVM.cpp diff --git a/third_party/mthreads/lib/Conversion/TritonGPUToLLVM/GatherOpToLLVM.cpp b/third_party/mthreads/spec_cpp/lib/Conversion/TritonGPUToLLVM/GatherOpToLLVM.cpp similarity index 100% rename from third_party/mthreads/lib/Conversion/TritonGPUToLLVM/GatherOpToLLVM.cpp rename to third_party/mthreads/spec_cpp/lib/Conversion/TritonGPUToLLVM/GatherOpToLLVM.cpp diff --git a/third_party/mthreads/lib/Conversion/TritonGPUToLLVM/GlobalScratchMemoryAllocation.cpp b/third_party/mthreads/spec_cpp/lib/Conversion/TritonGPUToLLVM/GlobalScratchMemoryAllocation.cpp similarity index 100% rename from third_party/mthreads/lib/Conversion/TritonGPUToLLVM/GlobalScratchMemoryAllocation.cpp rename to third_party/mthreads/spec_cpp/lib/Conversion/TritonGPUToLLVM/GlobalScratchMemoryAllocation.cpp diff --git a/third_party/mthreads/lib/Conversion/TritonGPUToLLVM/HistogramOpToLLVM.cpp b/third_party/mthreads/spec_cpp/lib/Conversion/TritonGPUToLLVM/HistogramOpToLLVM.cpp similarity index 100% rename from third_party/mthreads/lib/Conversion/TritonGPUToLLVM/HistogramOpToLLVM.cpp rename to third_party/mthreads/spec_cpp/lib/Conversion/TritonGPUToLLVM/HistogramOpToLLVM.cpp diff --git a/third_party/mthreads/lib/Conversion/TritonGPUToLLVM/MemoryOpToLLVM.cpp b/third_party/mthreads/spec_cpp/lib/Conversion/TritonGPUToLLVM/MemoryOpToLLVM.cpp similarity index 100% rename from third_party/mthreads/lib/Conversion/TritonGPUToLLVM/MemoryOpToLLVM.cpp rename to third_party/mthreads/spec_cpp/lib/Conversion/TritonGPUToLLVM/MemoryOpToLLVM.cpp diff --git a/third_party/mthreads/lib/Conversion/TritonGPUToLLVM/ReduceOpToLLVM.cpp b/third_party/mthreads/spec_cpp/lib/Conversion/TritonGPUToLLVM/ReduceOpToLLVM.cpp similarity index 100% rename from third_party/mthreads/lib/Conversion/TritonGPUToLLVM/ReduceOpToLLVM.cpp rename to third_party/mthreads/spec_cpp/lib/Conversion/TritonGPUToLLVM/ReduceOpToLLVM.cpp diff --git a/third_party/mthreads/lib/Conversion/TritonGPUToLLVM/ScanOpToLLVM.cpp b/third_party/mthreads/spec_cpp/lib/Conversion/TritonGPUToLLVM/ScanOpToLLVM.cpp similarity index 100% rename from third_party/mthreads/lib/Conversion/TritonGPUToLLVM/ScanOpToLLVM.cpp rename to third_party/mthreads/spec_cpp/lib/Conversion/TritonGPUToLLVM/ScanOpToLLVM.cpp diff --git a/third_party/mthreads/lib/Conversion/TritonGPUToLLVM/Utility.cpp b/third_party/mthreads/spec_cpp/lib/Conversion/TritonGPUToLLVM/Utility.cpp similarity index 100% rename from third_party/mthreads/lib/Conversion/TritonGPUToLLVM/Utility.cpp rename to third_party/mthreads/spec_cpp/lib/Conversion/TritonGPUToLLVM/Utility.cpp diff --git a/third_party/mthreads/lib/Conversion/TritonGPUToLLVM/ViewOpToLLVM.cpp b/third_party/mthreads/spec_cpp/lib/Conversion/TritonGPUToLLVM/ViewOpToLLVM.cpp similarity index 100% rename from third_party/mthreads/lib/Conversion/TritonGPUToLLVM/ViewOpToLLVM.cpp rename to third_party/mthreads/spec_cpp/lib/Conversion/TritonGPUToLLVM/ViewOpToLLVM.cpp diff --git a/third_party/mthreads/lib/Conversion/TritonGPUToLLVM/WarpSpecializeUtility.cpp b/third_party/mthreads/spec_cpp/lib/Conversion/TritonGPUToLLVM/WarpSpecializeUtility.cpp similarity index 100% rename from third_party/mthreads/lib/Conversion/TritonGPUToLLVM/WarpSpecializeUtility.cpp rename to third_party/mthreads/spec_cpp/lib/Conversion/TritonGPUToLLVM/WarpSpecializeUtility.cpp diff --git a/third_party/mthreads/lib/Conversion/TritonInstrumentToLLVM/InstrumentationToLLVM.cpp b/third_party/mthreads/spec_cpp/lib/Conversion/TritonInstrumentToLLVM/InstrumentationToLLVM.cpp similarity index 100% rename from third_party/mthreads/lib/Conversion/TritonInstrumentToLLVM/InstrumentationToLLVM.cpp rename to third_party/mthreads/spec_cpp/lib/Conversion/TritonInstrumentToLLVM/InstrumentationToLLVM.cpp diff --git a/third_party/mthreads/lib/Conversion/TritonToTritonGPU/RelayoutTritonGPU.cpp b/third_party/mthreads/spec_cpp/lib/Conversion/TritonToTritonGPU/RelayoutTritonGPU.cpp similarity index 100% rename from third_party/mthreads/lib/Conversion/TritonToTritonGPU/RelayoutTritonGPU.cpp rename to third_party/mthreads/spec_cpp/lib/Conversion/TritonToTritonGPU/RelayoutTritonGPU.cpp diff --git a/third_party/mthreads/lib/Conversion/TritonToTritonGPU/TritonGPUConversion.cpp b/third_party/mthreads/spec_cpp/lib/Conversion/TritonToTritonGPU/TritonGPUConversion.cpp similarity index 100% rename from third_party/mthreads/lib/Conversion/TritonToTritonGPU/TritonGPUConversion.cpp rename to third_party/mthreads/spec_cpp/lib/Conversion/TritonToTritonGPU/TritonGPUConversion.cpp diff --git a/third_party/mthreads/lib/Conversion/TritonToTritonGPU/TritonToTritonGPUPass.cpp b/third_party/mthreads/spec_cpp/lib/Conversion/TritonToTritonGPU/TritonToTritonGPUPass.cpp similarity index 100% rename from third_party/mthreads/lib/Conversion/TritonToTritonGPU/TritonToTritonGPUPass.cpp rename to third_party/mthreads/spec_cpp/lib/Conversion/TritonToTritonGPU/TritonToTritonGPUPass.cpp diff --git a/third_party/mthreads/lib/Dialect/Gluon/Transforms/Canonicalize.cpp b/third_party/mthreads/spec_cpp/lib/Dialect/Gluon/Transforms/Canonicalize.cpp similarity index 100% rename from third_party/mthreads/lib/Dialect/Gluon/Transforms/Canonicalize.cpp rename to third_party/mthreads/spec_cpp/lib/Dialect/Gluon/Transforms/Canonicalize.cpp diff --git a/third_party/mthreads/lib/Dialect/Gluon/Transforms/InferCoalescedEncodings.cpp b/third_party/mthreads/spec_cpp/lib/Dialect/Gluon/Transforms/InferCoalescedEncodings.cpp similarity index 100% rename from third_party/mthreads/lib/Dialect/Gluon/Transforms/InferCoalescedEncodings.cpp rename to third_party/mthreads/spec_cpp/lib/Dialect/Gluon/Transforms/InferCoalescedEncodings.cpp diff --git a/third_party/mthreads/lib/Dialect/NVGPU/IR/CMakeLists.txt b/third_party/mthreads/spec_cpp/lib/Dialect/NVGPU/IR/CMakeLists.txt similarity index 100% rename from third_party/mthreads/lib/Dialect/NVGPU/IR/CMakeLists.txt rename to third_party/mthreads/spec_cpp/lib/Dialect/NVGPU/IR/CMakeLists.txt diff --git a/third_party/mthreads/lib/Dialect/NVGPU/IR/Dialect.cpp b/third_party/mthreads/spec_cpp/lib/Dialect/NVGPU/IR/Dialect.cpp similarity index 100% rename from third_party/mthreads/lib/Dialect/NVGPU/IR/Dialect.cpp rename to third_party/mthreads/spec_cpp/lib/Dialect/NVGPU/IR/Dialect.cpp diff --git a/third_party/mthreads/lib/Dialect/NVWS/IR/CMakeLists.txt b/third_party/mthreads/spec_cpp/lib/Dialect/NVWS/IR/CMakeLists.txt similarity index 100% rename from third_party/mthreads/lib/Dialect/NVWS/IR/CMakeLists.txt rename to third_party/mthreads/spec_cpp/lib/Dialect/NVWS/IR/CMakeLists.txt diff --git a/third_party/mthreads/lib/Dialect/NVWS/IR/Dialect.cpp b/third_party/mthreads/spec_cpp/lib/Dialect/NVWS/IR/Dialect.cpp similarity index 100% rename from third_party/mthreads/lib/Dialect/NVWS/IR/Dialect.cpp rename to third_party/mthreads/spec_cpp/lib/Dialect/NVWS/IR/Dialect.cpp diff --git a/third_party/mthreads/lib/Dialect/NVWS/IR/Ops.cpp b/third_party/mthreads/spec_cpp/lib/Dialect/NVWS/IR/Ops.cpp similarity index 100% rename from third_party/mthreads/lib/Dialect/NVWS/IR/Ops.cpp rename to third_party/mthreads/spec_cpp/lib/Dialect/NVWS/IR/Ops.cpp diff --git a/third_party/mthreads/lib/Dialect/NVWS/Transforms/AssignStagePhase.cpp b/third_party/mthreads/spec_cpp/lib/Dialect/NVWS/Transforms/AssignStagePhase.cpp similarity index 100% rename from third_party/mthreads/lib/Dialect/NVWS/Transforms/AssignStagePhase.cpp rename to third_party/mthreads/spec_cpp/lib/Dialect/NVWS/Transforms/AssignStagePhase.cpp diff --git a/third_party/mthreads/lib/Dialect/NVWS/Transforms/CMakeLists.txt b/third_party/mthreads/spec_cpp/lib/Dialect/NVWS/Transforms/CMakeLists.txt similarity index 100% rename from third_party/mthreads/lib/Dialect/NVWS/Transforms/CMakeLists.txt rename to third_party/mthreads/spec_cpp/lib/Dialect/NVWS/Transforms/CMakeLists.txt diff --git a/third_party/mthreads/lib/Dialect/NVWS/Transforms/HoistTmemStore.cpp b/third_party/mthreads/spec_cpp/lib/Dialect/NVWS/Transforms/HoistTmemStore.cpp similarity index 100% rename from third_party/mthreads/lib/Dialect/NVWS/Transforms/HoistTmemStore.cpp rename to third_party/mthreads/spec_cpp/lib/Dialect/NVWS/Transforms/HoistTmemStore.cpp diff --git a/third_party/mthreads/lib/Dialect/NVWS/Transforms/InsertAref.cpp b/third_party/mthreads/spec_cpp/lib/Dialect/NVWS/Transforms/InsertAref.cpp similarity index 100% rename from third_party/mthreads/lib/Dialect/NVWS/Transforms/InsertAref.cpp rename to third_party/mthreads/spec_cpp/lib/Dialect/NVWS/Transforms/InsertAref.cpp diff --git a/third_party/mthreads/lib/Dialect/NVWS/Transforms/InsertTmemAref.cpp b/third_party/mthreads/spec_cpp/lib/Dialect/NVWS/Transforms/InsertTmemAref.cpp similarity index 100% rename from third_party/mthreads/lib/Dialect/NVWS/Transforms/InsertTmemAref.cpp rename to third_party/mthreads/spec_cpp/lib/Dialect/NVWS/Transforms/InsertTmemAref.cpp diff --git a/third_party/mthreads/lib/Dialect/NVWS/Transforms/LowerAref.cpp b/third_party/mthreads/spec_cpp/lib/Dialect/NVWS/Transforms/LowerAref.cpp similarity index 100% rename from third_party/mthreads/lib/Dialect/NVWS/Transforms/LowerAref.cpp rename to third_party/mthreads/spec_cpp/lib/Dialect/NVWS/Transforms/LowerAref.cpp diff --git a/third_party/mthreads/lib/Dialect/NVWS/Transforms/LowerWarpGroup.cpp b/third_party/mthreads/spec_cpp/lib/Dialect/NVWS/Transforms/LowerWarpGroup.cpp similarity index 100% rename from third_party/mthreads/lib/Dialect/NVWS/Transforms/LowerWarpGroup.cpp rename to third_party/mthreads/spec_cpp/lib/Dialect/NVWS/Transforms/LowerWarpGroup.cpp diff --git a/third_party/mthreads/lib/Dialect/NVWS/Transforms/Utilities.cpp b/third_party/mthreads/spec_cpp/lib/Dialect/NVWS/Transforms/Utilities.cpp similarity index 100% rename from third_party/mthreads/lib/Dialect/NVWS/Transforms/Utilities.cpp rename to third_party/mthreads/spec_cpp/lib/Dialect/NVWS/Transforms/Utilities.cpp diff --git a/third_party/mthreads/lib/Dialect/NVWS/Transforms/Utilities.h b/third_party/mthreads/spec_cpp/lib/Dialect/NVWS/Transforms/Utilities.h similarity index 100% rename from third_party/mthreads/lib/Dialect/NVWS/Transforms/Utilities.h rename to third_party/mthreads/spec_cpp/lib/Dialect/NVWS/Transforms/Utilities.h diff --git a/third_party/mthreads/lib/Dialect/Triton/IR/Ops.cpp b/third_party/mthreads/spec_cpp/lib/Dialect/Triton/IR/Ops.cpp similarity index 100% rename from third_party/mthreads/lib/Dialect/Triton/IR/Ops.cpp rename to third_party/mthreads/spec_cpp/lib/Dialect/Triton/IR/Ops.cpp diff --git a/third_party/mthreads/lib/Dialect/Triton/IR/Traits.cpp b/third_party/mthreads/spec_cpp/lib/Dialect/Triton/IR/Traits.cpp similarity index 100% rename from third_party/mthreads/lib/Dialect/Triton/IR/Traits.cpp rename to third_party/mthreads/spec_cpp/lib/Dialect/Triton/IR/Traits.cpp diff --git a/third_party/mthreads/lib/Dialect/Triton/IR/Types.cpp b/third_party/mthreads/spec_cpp/lib/Dialect/Triton/IR/Types.cpp similarity index 100% rename from third_party/mthreads/lib/Dialect/Triton/IR/Types.cpp rename to third_party/mthreads/spec_cpp/lib/Dialect/Triton/IR/Types.cpp diff --git a/third_party/mthreads/lib/Dialect/Triton/IR/Utility.cpp b/third_party/mthreads/spec_cpp/lib/Dialect/Triton/IR/Utility.cpp similarity index 100% rename from third_party/mthreads/lib/Dialect/Triton/IR/Utility.cpp rename to third_party/mthreads/spec_cpp/lib/Dialect/Triton/IR/Utility.cpp diff --git a/third_party/mthreads/lib/Dialect/Triton/Transforms/Combine.cpp b/third_party/mthreads/spec_cpp/lib/Dialect/Triton/Transforms/Combine.cpp similarity index 100% rename from third_party/mthreads/lib/Dialect/Triton/Transforms/Combine.cpp rename to third_party/mthreads/spec_cpp/lib/Dialect/Triton/Transforms/Combine.cpp diff --git a/third_party/mthreads/lib/Dialect/Triton/Transforms/Combine.td b/third_party/mthreads/spec_cpp/lib/Dialect/Triton/Transforms/Combine.td similarity index 100% rename from third_party/mthreads/lib/Dialect/Triton/Transforms/Combine.td rename to third_party/mthreads/spec_cpp/lib/Dialect/Triton/Transforms/Combine.td diff --git a/third_party/mthreads/lib/Dialect/Triton/Transforms/LoopUnroll.cpp b/third_party/mthreads/spec_cpp/lib/Dialect/Triton/Transforms/LoopUnroll.cpp similarity index 100% rename from third_party/mthreads/lib/Dialect/Triton/Transforms/LoopUnroll.cpp rename to third_party/mthreads/spec_cpp/lib/Dialect/Triton/Transforms/LoopUnroll.cpp diff --git a/third_party/mthreads/lib/Dialect/Triton/Transforms/RewriteTensorDescriptorToPointer.cpp b/third_party/mthreads/spec_cpp/lib/Dialect/Triton/Transforms/RewriteTensorDescriptorToPointer.cpp similarity index 100% rename from third_party/mthreads/lib/Dialect/Triton/Transforms/RewriteTensorDescriptorToPointer.cpp rename to third_party/mthreads/spec_cpp/lib/Dialect/Triton/Transforms/RewriteTensorDescriptorToPointer.cpp diff --git a/third_party/mthreads/lib/Dialect/Triton/Transforms/RewriteTensorPointer.cpp b/third_party/mthreads/spec_cpp/lib/Dialect/Triton/Transforms/RewriteTensorPointer.cpp similarity index 100% rename from third_party/mthreads/lib/Dialect/Triton/Transforms/RewriteTensorPointer.cpp rename to third_party/mthreads/spec_cpp/lib/Dialect/Triton/Transforms/RewriteTensorPointer.cpp diff --git a/third_party/mthreads/lib/Dialect/TritonGPU/IR/Dialect.cpp b/third_party/mthreads/spec_cpp/lib/Dialect/TritonGPU/IR/Dialect.cpp similarity index 100% rename from third_party/mthreads/lib/Dialect/TritonGPU/IR/Dialect.cpp rename to third_party/mthreads/spec_cpp/lib/Dialect/TritonGPU/IR/Dialect.cpp diff --git a/third_party/mthreads/lib/Dialect/TritonGPU/IR/LinearLayoutConversions.cpp b/third_party/mthreads/spec_cpp/lib/Dialect/TritonGPU/IR/LinearLayoutConversions.cpp similarity index 100% rename from third_party/mthreads/lib/Dialect/TritonGPU/IR/LinearLayoutConversions.cpp rename to third_party/mthreads/spec_cpp/lib/Dialect/TritonGPU/IR/LinearLayoutConversions.cpp diff --git a/third_party/mthreads/lib/Dialect/TritonGPU/IR/Ops.cpp b/third_party/mthreads/spec_cpp/lib/Dialect/TritonGPU/IR/Ops.cpp similarity index 100% rename from third_party/mthreads/lib/Dialect/TritonGPU/IR/Ops.cpp rename to third_party/mthreads/spec_cpp/lib/Dialect/TritonGPU/IR/Ops.cpp diff --git a/third_party/mthreads/lib/Dialect/TritonGPU/IR/Types.cpp b/third_party/mthreads/spec_cpp/lib/Dialect/TritonGPU/IR/Types.cpp similarity index 100% rename from third_party/mthreads/lib/Dialect/TritonGPU/IR/Types.cpp rename to third_party/mthreads/spec_cpp/lib/Dialect/TritonGPU/IR/Types.cpp diff --git a/third_party/mthreads/lib/Dialect/TritonGPU/Transforms/AccelerateMatmul.cpp b/third_party/mthreads/spec_cpp/lib/Dialect/TritonGPU/Transforms/AccelerateMatmul.cpp similarity index 100% rename from third_party/mthreads/lib/Dialect/TritonGPU/Transforms/AccelerateMatmul.cpp rename to third_party/mthreads/spec_cpp/lib/Dialect/TritonGPU/Transforms/AccelerateMatmul.cpp diff --git a/third_party/mthreads/lib/Dialect/TritonGPU/Transforms/Coalesce.cpp b/third_party/mthreads/spec_cpp/lib/Dialect/TritonGPU/Transforms/Coalesce.cpp similarity index 100% rename from third_party/mthreads/lib/Dialect/TritonGPU/Transforms/Coalesce.cpp rename to third_party/mthreads/spec_cpp/lib/Dialect/TritonGPU/Transforms/Coalesce.cpp diff --git a/third_party/mthreads/lib/Dialect/TritonGPU/Transforms/CoalesceAsyncCopy.cpp b/third_party/mthreads/spec_cpp/lib/Dialect/TritonGPU/Transforms/CoalesceAsyncCopy.cpp similarity index 100% rename from third_party/mthreads/lib/Dialect/TritonGPU/Transforms/CoalesceAsyncCopy.cpp rename to third_party/mthreads/spec_cpp/lib/Dialect/TritonGPU/Transforms/CoalesceAsyncCopy.cpp diff --git a/third_party/mthreads/lib/Dialect/TritonGPU/Transforms/CoalesceUtils.cpp b/third_party/mthreads/spec_cpp/lib/Dialect/TritonGPU/Transforms/CoalesceUtils.cpp similarity index 100% rename from third_party/mthreads/lib/Dialect/TritonGPU/Transforms/CoalesceUtils.cpp rename to third_party/mthreads/spec_cpp/lib/Dialect/TritonGPU/Transforms/CoalesceUtils.cpp diff --git a/third_party/mthreads/lib/Dialect/TritonGPU/Transforms/F32DotTC.cpp b/third_party/mthreads/spec_cpp/lib/Dialect/TritonGPU/Transforms/F32DotTC.cpp similarity index 100% rename from third_party/mthreads/lib/Dialect/TritonGPU/Transforms/F32DotTC.cpp rename to third_party/mthreads/spec_cpp/lib/Dialect/TritonGPU/Transforms/F32DotTC.cpp diff --git a/third_party/mthreads/lib/Dialect/TritonGPU/Transforms/OptimizeAccumulatorInit.cpp b/third_party/mthreads/spec_cpp/lib/Dialect/TritonGPU/Transforms/OptimizeAccumulatorInit.cpp similarity index 100% rename from third_party/mthreads/lib/Dialect/TritonGPU/Transforms/OptimizeAccumulatorInit.cpp rename to third_party/mthreads/spec_cpp/lib/Dialect/TritonGPU/Transforms/OptimizeAccumulatorInit.cpp diff --git a/third_party/mthreads/lib/Dialect/TritonGPU/Transforms/OptimizeDotOperands.cpp b/third_party/mthreads/spec_cpp/lib/Dialect/TritonGPU/Transforms/OptimizeDotOperands.cpp similarity index 100% rename from third_party/mthreads/lib/Dialect/TritonGPU/Transforms/OptimizeDotOperands.cpp rename to third_party/mthreads/spec_cpp/lib/Dialect/TritonGPU/Transforms/OptimizeDotOperands.cpp diff --git a/third_party/mthreads/lib/Dialect/TritonGPU/Transforms/OptimizeThreadLocality.cpp b/third_party/mthreads/spec_cpp/lib/Dialect/TritonGPU/Transforms/OptimizeThreadLocality.cpp similarity index 100% rename from third_party/mthreads/lib/Dialect/TritonGPU/Transforms/OptimizeThreadLocality.cpp rename to third_party/mthreads/spec_cpp/lib/Dialect/TritonGPU/Transforms/OptimizeThreadLocality.cpp diff --git a/third_party/mthreads/lib/Dialect/TritonGPU/Transforms/Pipeliner/AssignLatencies.cpp b/third_party/mthreads/spec_cpp/lib/Dialect/TritonGPU/Transforms/Pipeliner/AssignLatencies.cpp similarity index 100% rename from third_party/mthreads/lib/Dialect/TritonGPU/Transforms/Pipeliner/AssignLatencies.cpp rename to third_party/mthreads/spec_cpp/lib/Dialect/TritonGPU/Transforms/Pipeliner/AssignLatencies.cpp diff --git a/third_party/mthreads/lib/Dialect/TritonGPU/Transforms/Pipeliner/LowerLoops.cpp b/third_party/mthreads/spec_cpp/lib/Dialect/TritonGPU/Transforms/Pipeliner/LowerLoops.cpp similarity index 100% rename from third_party/mthreads/lib/Dialect/TritonGPU/Transforms/Pipeliner/LowerLoops.cpp rename to third_party/mthreads/spec_cpp/lib/Dialect/TritonGPU/Transforms/Pipeliner/LowerLoops.cpp diff --git a/third_party/mthreads/lib/Dialect/TritonGPU/Transforms/Pipeliner/MMAv5PipelineUtility.cpp b/third_party/mthreads/spec_cpp/lib/Dialect/TritonGPU/Transforms/Pipeliner/MMAv5PipelineUtility.cpp similarity index 100% rename from third_party/mthreads/lib/Dialect/TritonGPU/Transforms/Pipeliner/MMAv5PipelineUtility.cpp rename to third_party/mthreads/spec_cpp/lib/Dialect/TritonGPU/Transforms/Pipeliner/MMAv5PipelineUtility.cpp diff --git a/third_party/mthreads/lib/Dialect/TritonGPU/Transforms/Pipeliner/PipelineExpander.cpp b/third_party/mthreads/spec_cpp/lib/Dialect/TritonGPU/Transforms/Pipeliner/PipelineExpander.cpp similarity index 100% rename from third_party/mthreads/lib/Dialect/TritonGPU/Transforms/Pipeliner/PipelineExpander.cpp rename to third_party/mthreads/spec_cpp/lib/Dialect/TritonGPU/Transforms/Pipeliner/PipelineExpander.cpp diff --git a/third_party/mthreads/lib/Dialect/TritonGPU/Transforms/Pipeliner/PipeliningUtility.cpp b/third_party/mthreads/spec_cpp/lib/Dialect/TritonGPU/Transforms/Pipeliner/PipeliningUtility.cpp similarity index 100% rename from third_party/mthreads/lib/Dialect/TritonGPU/Transforms/Pipeliner/PipeliningUtility.cpp rename to third_party/mthreads/spec_cpp/lib/Dialect/TritonGPU/Transforms/Pipeliner/PipeliningUtility.cpp diff --git a/third_party/mthreads/lib/Dialect/TritonGPU/Transforms/Pipeliner/Schedule.cpp b/third_party/mthreads/spec_cpp/lib/Dialect/TritonGPU/Transforms/Pipeliner/Schedule.cpp similarity index 100% rename from third_party/mthreads/lib/Dialect/TritonGPU/Transforms/Pipeliner/Schedule.cpp rename to third_party/mthreads/spec_cpp/lib/Dialect/TritonGPU/Transforms/Pipeliner/Schedule.cpp diff --git a/third_party/mthreads/lib/Dialect/TritonGPU/Transforms/Pipeliner/ScheduleLoops.cpp b/third_party/mthreads/spec_cpp/lib/Dialect/TritonGPU/Transforms/Pipeliner/ScheduleLoops.cpp similarity index 100% rename from third_party/mthreads/lib/Dialect/TritonGPU/Transforms/Pipeliner/ScheduleLoops.cpp rename to third_party/mthreads/spec_cpp/lib/Dialect/TritonGPU/Transforms/Pipeliner/ScheduleLoops.cpp diff --git a/third_party/mthreads/lib/Dialect/TritonGPU/Transforms/Pipeliner/SoftwarePipeliner.cpp b/third_party/mthreads/spec_cpp/lib/Dialect/TritonGPU/Transforms/Pipeliner/SoftwarePipeliner.cpp similarity index 100% rename from third_party/mthreads/lib/Dialect/TritonGPU/Transforms/Pipeliner/SoftwarePipeliner.cpp rename to third_party/mthreads/spec_cpp/lib/Dialect/TritonGPU/Transforms/Pipeliner/SoftwarePipeliner.cpp diff --git a/third_party/mthreads/lib/Dialect/TritonGPU/Transforms/Pipeliner/TMAStoresPipeline.cpp b/third_party/mthreads/spec_cpp/lib/Dialect/TritonGPU/Transforms/Pipeliner/TMAStoresPipeline.cpp similarity index 100% rename from third_party/mthreads/lib/Dialect/TritonGPU/Transforms/Pipeliner/TMAStoresPipeline.cpp rename to third_party/mthreads/spec_cpp/lib/Dialect/TritonGPU/Transforms/Pipeliner/TMAStoresPipeline.cpp diff --git a/third_party/mthreads/lib/Dialect/TritonGPU/Transforms/Pipeliner/WGMMAPipeline.cpp b/third_party/mthreads/spec_cpp/lib/Dialect/TritonGPU/Transforms/Pipeliner/WGMMAPipeline.cpp similarity index 100% rename from third_party/mthreads/lib/Dialect/TritonGPU/Transforms/Pipeliner/WGMMAPipeline.cpp rename to third_party/mthreads/spec_cpp/lib/Dialect/TritonGPU/Transforms/Pipeliner/WGMMAPipeline.cpp diff --git a/third_party/mthreads/lib/Dialect/TritonGPU/Transforms/Prefetch.cpp b/third_party/mthreads/spec_cpp/lib/Dialect/TritonGPU/Transforms/Prefetch.cpp similarity index 100% rename from third_party/mthreads/lib/Dialect/TritonGPU/Transforms/Prefetch.cpp rename to third_party/mthreads/spec_cpp/lib/Dialect/TritonGPU/Transforms/Prefetch.cpp diff --git a/third_party/mthreads/lib/Dialect/TritonGPU/Transforms/ReduceDataDuplication.cpp b/third_party/mthreads/spec_cpp/lib/Dialect/TritonGPU/Transforms/ReduceDataDuplication.cpp similarity index 100% rename from third_party/mthreads/lib/Dialect/TritonGPU/Transforms/ReduceDataDuplication.cpp rename to third_party/mthreads/spec_cpp/lib/Dialect/TritonGPU/Transforms/ReduceDataDuplication.cpp diff --git a/third_party/mthreads/lib/Dialect/TritonGPU/Transforms/RemoveLayoutConversions.cpp b/third_party/mthreads/spec_cpp/lib/Dialect/TritonGPU/Transforms/RemoveLayoutConversions.cpp similarity index 100% rename from third_party/mthreads/lib/Dialect/TritonGPU/Transforms/RemoveLayoutConversions.cpp rename to third_party/mthreads/spec_cpp/lib/Dialect/TritonGPU/Transforms/RemoveLayoutConversions.cpp diff --git a/third_party/mthreads/lib/Dialect/TritonGPU/Transforms/Utility.cpp b/third_party/mthreads/spec_cpp/lib/Dialect/TritonGPU/Transforms/Utility.cpp similarity index 100% rename from third_party/mthreads/lib/Dialect/TritonGPU/Transforms/Utility.cpp rename to third_party/mthreads/spec_cpp/lib/Dialect/TritonGPU/Transforms/Utility.cpp diff --git a/third_party/mthreads/lib/Dialect/TritonGPU/Transforms/WarpSpecialization/AutomaticWarpSpecialization.cpp b/third_party/mthreads/spec_cpp/lib/Dialect/TritonGPU/Transforms/WarpSpecialization/AutomaticWarpSpecialization.cpp similarity index 100% rename from third_party/mthreads/lib/Dialect/TritonGPU/Transforms/WarpSpecialization/AutomaticWarpSpecialization.cpp rename to third_party/mthreads/spec_cpp/lib/Dialect/TritonGPU/Transforms/WarpSpecialization/AutomaticWarpSpecialization.cpp diff --git a/third_party/mthreads/lib/Dialect/TritonGPU/Transforms/WarpSpecialization/OptimizePartitionWarps.cpp b/third_party/mthreads/spec_cpp/lib/Dialect/TritonGPU/Transforms/WarpSpecialization/OptimizePartitionWarps.cpp similarity index 100% rename from third_party/mthreads/lib/Dialect/TritonGPU/Transforms/WarpSpecialization/OptimizePartitionWarps.cpp rename to third_party/mthreads/spec_cpp/lib/Dialect/TritonGPU/Transforms/WarpSpecialization/OptimizePartitionWarps.cpp diff --git a/third_party/mthreads/lib/Dialect/TritonGPU/Transforms/WarpSpecialization/Partition.cpp b/third_party/mthreads/spec_cpp/lib/Dialect/TritonGPU/Transforms/WarpSpecialization/Partition.cpp similarity index 100% rename from third_party/mthreads/lib/Dialect/TritonGPU/Transforms/WarpSpecialization/Partition.cpp rename to third_party/mthreads/spec_cpp/lib/Dialect/TritonGPU/Transforms/WarpSpecialization/Partition.cpp diff --git a/third_party/mthreads/lib/Dialect/TritonGPU/Transforms/WarpSpecialization/PartitionLoops.cpp b/third_party/mthreads/spec_cpp/lib/Dialect/TritonGPU/Transforms/WarpSpecialization/PartitionLoops.cpp similarity index 100% rename from third_party/mthreads/lib/Dialect/TritonGPU/Transforms/WarpSpecialization/PartitionLoops.cpp rename to third_party/mthreads/spec_cpp/lib/Dialect/TritonGPU/Transforms/WarpSpecialization/PartitionLoops.cpp diff --git a/third_party/mthreads/lib/Dialect/TritonGPU/Transforms/WarpSpecialization/PartitionScheduling.cpp b/third_party/mthreads/spec_cpp/lib/Dialect/TritonGPU/Transforms/WarpSpecialization/PartitionScheduling.cpp similarity index 100% rename from third_party/mthreads/lib/Dialect/TritonGPU/Transforms/WarpSpecialization/PartitionScheduling.cpp rename to third_party/mthreads/spec_cpp/lib/Dialect/TritonGPU/Transforms/WarpSpecialization/PartitionScheduling.cpp diff --git a/third_party/mthreads/lib/Dialect/TritonGPU/Transforms/WarpSpecialization/PartitionSchedulingUtility.cpp b/third_party/mthreads/spec_cpp/lib/Dialect/TritonGPU/Transforms/WarpSpecialization/PartitionSchedulingUtility.cpp similarity index 100% rename from third_party/mthreads/lib/Dialect/TritonGPU/Transforms/WarpSpecialization/PartitionSchedulingUtility.cpp rename to third_party/mthreads/spec_cpp/lib/Dialect/TritonGPU/Transforms/WarpSpecialization/PartitionSchedulingUtility.cpp diff --git a/third_party/mthreads/lib/Dialect/TritonInstrument/IR/FunctionBuilder.cpp b/third_party/mthreads/spec_cpp/lib/Dialect/TritonInstrument/IR/FunctionBuilder.cpp similarity index 100% rename from third_party/mthreads/lib/Dialect/TritonInstrument/IR/FunctionBuilder.cpp rename to third_party/mthreads/spec_cpp/lib/Dialect/TritonInstrument/IR/FunctionBuilder.cpp diff --git a/third_party/mthreads/lib/Dialect/TritonInstrument/IR/Utility.cpp b/third_party/mthreads/spec_cpp/lib/Dialect/TritonInstrument/IR/Utility.cpp similarity index 100% rename from third_party/mthreads/lib/Dialect/TritonInstrument/IR/Utility.cpp rename to third_party/mthreads/spec_cpp/lib/Dialect/TritonInstrument/IR/Utility.cpp diff --git a/third_party/mthreads/lib/Dialect/TritonInstrument/Transforms/ConcurrencySanitizer.cpp b/third_party/mthreads/spec_cpp/lib/Dialect/TritonInstrument/Transforms/ConcurrencySanitizer.cpp similarity index 100% rename from third_party/mthreads/lib/Dialect/TritonInstrument/Transforms/ConcurrencySanitizer.cpp rename to third_party/mthreads/spec_cpp/lib/Dialect/TritonInstrument/Transforms/ConcurrencySanitizer.cpp diff --git a/third_party/mthreads/lib/Dialect/TritonNvidiaGPU/IR/Dialect.cpp b/third_party/mthreads/spec_cpp/lib/Dialect/TritonNvidiaGPU/IR/Dialect.cpp similarity index 100% rename from third_party/mthreads/lib/Dialect/TritonNvidiaGPU/IR/Dialect.cpp rename to third_party/mthreads/spec_cpp/lib/Dialect/TritonNvidiaGPU/IR/Dialect.cpp diff --git a/third_party/mthreads/lib/Dialect/TritonNvidiaGPU/IR/Ops.cpp b/third_party/mthreads/spec_cpp/lib/Dialect/TritonNvidiaGPU/IR/Ops.cpp similarity index 100% rename from third_party/mthreads/lib/Dialect/TritonNvidiaGPU/IR/Ops.cpp rename to third_party/mthreads/spec_cpp/lib/Dialect/TritonNvidiaGPU/IR/Ops.cpp diff --git a/third_party/mthreads/lib/Dialect/TritonNvidiaGPU/IR/TensorMemoryUtils.cpp b/third_party/mthreads/spec_cpp/lib/Dialect/TritonNvidiaGPU/IR/TensorMemoryUtils.cpp similarity index 100% rename from third_party/mthreads/lib/Dialect/TritonNvidiaGPU/IR/TensorMemoryUtils.cpp rename to third_party/mthreads/spec_cpp/lib/Dialect/TritonNvidiaGPU/IR/TensorMemoryUtils.cpp diff --git a/third_party/mthreads/lib/Dialect/TritonNvidiaGPU/Transforms/FenceInsertion.cpp b/third_party/mthreads/spec_cpp/lib/Dialect/TritonNvidiaGPU/Transforms/FenceInsertion.cpp similarity index 100% rename from third_party/mthreads/lib/Dialect/TritonNvidiaGPU/Transforms/FenceInsertion.cpp rename to third_party/mthreads/spec_cpp/lib/Dialect/TritonNvidiaGPU/Transforms/FenceInsertion.cpp diff --git a/third_party/mthreads/lib/Dialect/TritonNvidiaGPU/Transforms/InterleaveTMem.cpp b/third_party/mthreads/spec_cpp/lib/Dialect/TritonNvidiaGPU/Transforms/InterleaveTMem.cpp similarity index 100% rename from third_party/mthreads/lib/Dialect/TritonNvidiaGPU/Transforms/InterleaveTMem.cpp rename to third_party/mthreads/spec_cpp/lib/Dialect/TritonNvidiaGPU/Transforms/InterleaveTMem.cpp diff --git a/third_party/mthreads/lib/Dialect/TritonNvidiaGPU/Transforms/MMALowering.cpp b/third_party/mthreads/spec_cpp/lib/Dialect/TritonNvidiaGPU/Transforms/MMALowering.cpp similarity index 100% rename from third_party/mthreads/lib/Dialect/TritonNvidiaGPU/Transforms/MMALowering.cpp rename to third_party/mthreads/spec_cpp/lib/Dialect/TritonNvidiaGPU/Transforms/MMALowering.cpp diff --git a/third_party/mthreads/lib/Dialect/TritonNvidiaGPU/Transforms/OptimizeDescriptorEncoding.cpp b/third_party/mthreads/spec_cpp/lib/Dialect/TritonNvidiaGPU/Transforms/OptimizeDescriptorEncoding.cpp similarity index 100% rename from third_party/mthreads/lib/Dialect/TritonNvidiaGPU/Transforms/OptimizeDescriptorEncoding.cpp rename to third_party/mthreads/spec_cpp/lib/Dialect/TritonNvidiaGPU/Transforms/OptimizeDescriptorEncoding.cpp diff --git a/third_party/mthreads/lib/Dialect/TritonNvidiaGPU/Transforms/OptimizeTMemLayouts.cpp b/third_party/mthreads/spec_cpp/lib/Dialect/TritonNvidiaGPU/Transforms/OptimizeTMemLayouts.cpp similarity index 100% rename from third_party/mthreads/lib/Dialect/TritonNvidiaGPU/Transforms/OptimizeTMemLayouts.cpp rename to third_party/mthreads/spec_cpp/lib/Dialect/TritonNvidiaGPU/Transforms/OptimizeTMemLayouts.cpp diff --git a/third_party/mthreads/lib/Dialect/TritonNvidiaGPU/Transforms/PlanCTA.cpp b/third_party/mthreads/spec_cpp/lib/Dialect/TritonNvidiaGPU/Transforms/PlanCTA.cpp similarity index 100% rename from third_party/mthreads/lib/Dialect/TritonNvidiaGPU/Transforms/PlanCTA.cpp rename to third_party/mthreads/spec_cpp/lib/Dialect/TritonNvidiaGPU/Transforms/PlanCTA.cpp diff --git a/third_party/mthreads/lib/Dialect/TritonNvidiaGPU/Transforms/PromoteLHSToTMem.cpp b/third_party/mthreads/spec_cpp/lib/Dialect/TritonNvidiaGPU/Transforms/PromoteLHSToTMem.cpp similarity index 100% rename from third_party/mthreads/lib/Dialect/TritonNvidiaGPU/Transforms/PromoteLHSToTMem.cpp rename to third_party/mthreads/spec_cpp/lib/Dialect/TritonNvidiaGPU/Transforms/PromoteLHSToTMem.cpp diff --git a/third_party/mthreads/lib/Dialect/TritonNvidiaGPU/Transforms/ProxyFenceInsertion.cpp b/third_party/mthreads/spec_cpp/lib/Dialect/TritonNvidiaGPU/Transforms/ProxyFenceInsertion.cpp similarity index 100% rename from third_party/mthreads/lib/Dialect/TritonNvidiaGPU/Transforms/ProxyFenceInsertion.cpp rename to third_party/mthreads/spec_cpp/lib/Dialect/TritonNvidiaGPU/Transforms/ProxyFenceInsertion.cpp diff --git a/third_party/mthreads/lib/Dialect/TritonNvidiaGPU/Transforms/RemoveTMEMTokens.cpp b/third_party/mthreads/spec_cpp/lib/Dialect/TritonNvidiaGPU/Transforms/RemoveTMEMTokens.cpp similarity index 100% rename from third_party/mthreads/lib/Dialect/TritonNvidiaGPU/Transforms/RemoveTMEMTokens.cpp rename to third_party/mthreads/spec_cpp/lib/Dialect/TritonNvidiaGPU/Transforms/RemoveTMEMTokens.cpp diff --git a/third_party/mthreads/lib/Dialect/TritonNvidiaGPU/Transforms/TMALowering.cpp b/third_party/mthreads/spec_cpp/lib/Dialect/TritonNvidiaGPU/Transforms/TMALowering.cpp similarity index 100% rename from third_party/mthreads/lib/Dialect/TritonNvidiaGPU/Transforms/TMALowering.cpp rename to third_party/mthreads/spec_cpp/lib/Dialect/TritonNvidiaGPU/Transforms/TMALowering.cpp diff --git a/third_party/mthreads/lib/Dialect/TritonNvidiaGPU/Transforms/TMAUtilities.cpp b/third_party/mthreads/spec_cpp/lib/Dialect/TritonNvidiaGPU/Transforms/TMAUtilities.cpp similarity index 100% rename from third_party/mthreads/lib/Dialect/TritonNvidiaGPU/Transforms/TMAUtilities.cpp rename to third_party/mthreads/spec_cpp/lib/Dialect/TritonNvidiaGPU/Transforms/TMAUtilities.cpp diff --git a/third_party/mthreads/lib/Dialect/TritonNvidiaGPU/Transforms/TensorMemoryAllocation.cpp b/third_party/mthreads/spec_cpp/lib/Dialect/TritonNvidiaGPU/Transforms/TensorMemoryAllocation.cpp similarity index 100% rename from third_party/mthreads/lib/Dialect/TritonNvidiaGPU/Transforms/TensorMemoryAllocation.cpp rename to third_party/mthreads/spec_cpp/lib/Dialect/TritonNvidiaGPU/Transforms/TensorMemoryAllocation.cpp diff --git a/third_party/mthreads/lib/Target/LLVMIR/LLVMDILocalVariable.cpp b/third_party/mthreads/spec_cpp/lib/Target/LLVMIR/LLVMDILocalVariable.cpp similarity index 100% rename from third_party/mthreads/lib/Target/LLVMIR/LLVMDILocalVariable.cpp rename to third_party/mthreads/spec_cpp/lib/Target/LLVMIR/LLVMDILocalVariable.cpp diff --git a/third_party/mthreads/lib/Target/LLVMIR/LLVMDIScope.cpp b/third_party/mthreads/spec_cpp/lib/Target/LLVMIR/LLVMDIScope.cpp similarity index 100% rename from third_party/mthreads/lib/Target/LLVMIR/LLVMDIScope.cpp rename to third_party/mthreads/spec_cpp/lib/Target/LLVMIR/LLVMDIScope.cpp diff --git a/third_party/mthreads/lib/Target/LLVMIR/LLVMDIUtils.cpp b/third_party/mthreads/spec_cpp/lib/Target/LLVMIR/LLVMDIUtils.cpp similarity index 100% rename from third_party/mthreads/lib/Target/LLVMIR/LLVMDIUtils.cpp rename to third_party/mthreads/spec_cpp/lib/Target/LLVMIR/LLVMDIUtils.cpp diff --git a/third_party/mthreads/lib/Target/LLVMIR/LLVMDIUtils.h b/third_party/mthreads/spec_cpp/lib/Target/LLVMIR/LLVMDIUtils.h similarity index 100% rename from third_party/mthreads/lib/Target/LLVMIR/LLVMDIUtils.h rename to third_party/mthreads/spec_cpp/lib/Target/LLVMIR/LLVMDIUtils.h diff --git a/third_party/mthreads/lib/Tools/PluginUtils.cpp b/third_party/mthreads/spec_cpp/lib/Tools/PluginUtils.cpp similarity index 100% rename from third_party/mthreads/lib/Tools/PluginUtils.cpp rename to third_party/mthreads/spec_cpp/lib/Tools/PluginUtils.cpp From 7a6f97c41309fe1c75e751b7b599ec9e5ad70115 Mon Sep 17 00:00:00 2001 From: Stardep <1486216685@qq.com> Date: Fri, 14 Aug 2026 09:50:02 +0800 Subject: [PATCH 2/4] [SPEC][MTHREADS] Fix Proton specialized header include --- third_party/mthreads/proton/Dialect/include/compat/Utility.h | 2 +- third_party/mthreads/proton/Dialect/lib/compat/TargetInfo.cpp | 2 +- third_party/mthreads/proton/Dialect/lib/compat/Utility.cpp | 2 +- 3 files changed, 3 insertions(+), 3 deletions(-) diff --git a/third_party/mthreads/proton/Dialect/include/compat/Utility.h b/third_party/mthreads/proton/Dialect/include/compat/Utility.h index 908540e9d0..bfe961469f 100644 --- a/third_party/mthreads/proton/Dialect/include/compat/Utility.h +++ b/third_party/mthreads/proton/Dialect/include/compat/Utility.h @@ -10,7 +10,7 @@ #include "compat/TargetInfo.h" #include "mlir/Conversion/LLVMCommon/Pattern.h" #include "mlir/Dialect/LLVMIR/LLVMDialect.h" -#include "third_party/mthreads/include/triton/Dialect/NVGPU/IR/Dialect.h" +#include "triton/Dialect/NVGPU/IR/Dialect.h" #include "triton/Analysis/Utility.h" #include "triton/Conversion/MLIRTypes.h" #include "triton/Dialect/TritonNvidiaGPU/IR/Dialect.h" diff --git a/third_party/mthreads/proton/Dialect/lib/compat/TargetInfo.cpp b/third_party/mthreads/proton/Dialect/lib/compat/TargetInfo.cpp index 2f39619c11..44bf46e18f 100644 --- a/third_party/mthreads/proton/Dialect/lib/compat/TargetInfo.cpp +++ b/third_party/mthreads/proton/Dialect/lib/compat/TargetInfo.cpp @@ -4,7 +4,7 @@ #include "mlir/Dialect/LLVMIR/LLVMDialect.h" #include "mlir/Dialect/LLVMIR/LLVMTypes.h" #include "mlir/Dialect/LLVMIR/NVVMDialect.h" -#include "third_party/mthreads/include/triton/Dialect/NVGPU/IR/Dialect.h" +#include "triton/Dialect/NVGPU/IR/Dialect.h" #include "triton/Dialect/TritonGPU/Transforms/Utility.h" #include "llvm/Support/MathExtras.h" diff --git a/third_party/mthreads/proton/Dialect/lib/compat/Utility.cpp b/third_party/mthreads/proton/Dialect/lib/compat/Utility.cpp index 95c2e179ee..cfd5d1fa95 100644 --- a/third_party/mthreads/proton/Dialect/lib/compat/Utility.cpp +++ b/third_party/mthreads/proton/Dialect/lib/compat/Utility.cpp @@ -1,6 +1,6 @@ #include "compat/Utility.h" #include "mlir/Dialect/LLVMIR/NVVMDialect.h" -#include "third_party/mthreads/include/triton/Dialect/NVGPU/IR/Dialect.h" +#include "triton/Dialect/NVGPU/IR/Dialect.h" #include "triton/Conversion/TritonGPUToLLVM/Utility.h" #include "triton/Tools/LayoutUtils.h" #include "triton/Tools/LinearLayout.h" From ddc4e4da8755683e32b849e0d3a2d756093bbec7 Mon Sep 17 00:00:00 2001 From: flagtree-bot Date: Fri, 14 Aug 2026 02:03:32 +0000 Subject: [PATCH 3/4] Apply code-format changes --- third_party/mthreads/proton/Dialect/include/compat/Utility.h | 2 +- third_party/mthreads/proton/Dialect/lib/compat/Utility.cpp | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/third_party/mthreads/proton/Dialect/include/compat/Utility.h b/third_party/mthreads/proton/Dialect/include/compat/Utility.h index bfe961469f..86b5d4b82f 100644 --- a/third_party/mthreads/proton/Dialect/include/compat/Utility.h +++ b/third_party/mthreads/proton/Dialect/include/compat/Utility.h @@ -10,9 +10,9 @@ #include "compat/TargetInfo.h" #include "mlir/Conversion/LLVMCommon/Pattern.h" #include "mlir/Dialect/LLVMIR/LLVMDialect.h" -#include "triton/Dialect/NVGPU/IR/Dialect.h" #include "triton/Analysis/Utility.h" #include "triton/Conversion/MLIRTypes.h" +#include "triton/Dialect/NVGPU/IR/Dialect.h" #include "triton/Dialect/TritonNvidiaGPU/IR/Dialect.h" #define DEBUG_TYPE "ttgpu_to_llvm" diff --git a/third_party/mthreads/proton/Dialect/lib/compat/Utility.cpp b/third_party/mthreads/proton/Dialect/lib/compat/Utility.cpp index cfd5d1fa95..e0a0f0c00c 100644 --- a/third_party/mthreads/proton/Dialect/lib/compat/Utility.cpp +++ b/third_party/mthreads/proton/Dialect/lib/compat/Utility.cpp @@ -1,7 +1,7 @@ #include "compat/Utility.h" #include "mlir/Dialect/LLVMIR/NVVMDialect.h" -#include "triton/Dialect/NVGPU/IR/Dialect.h" #include "triton/Conversion/TritonGPUToLLVM/Utility.h" +#include "triton/Dialect/NVGPU/IR/Dialect.h" #include "triton/Tools/LayoutUtils.h" #include "triton/Tools/LinearLayout.h" #include "llvm/Support/raw_ostream.h" From f737313519a5632e6c6242c4574eda14d7fdd257 Mon Sep 17 00:00:00 2001 From: Stardep <1486216685@qq.com> Date: Fri, 14 Aug 2026 10:18:19 +0800 Subject: [PATCH 4/4] [SPEC] Use backend TD files for MLIR docs --- cmake/FlagTreeOptions.cmake | 11 +++++++++++ include/triton/Dialect/Triton/IR/CMakeLists.txt | 2 +- include/triton/Dialect/TritonGPU/IR/CMakeLists.txt | 2 +- .../triton/Dialect/TritonInstrument/IR/CMakeLists.txt | 2 +- .../triton/Dialect/TritonNvidiaGPU/IR/CMakeLists.txt | 4 ++-- 5 files changed, 16 insertions(+), 5 deletions(-) diff --git a/cmake/FlagTreeOptions.cmake b/cmake/FlagTreeOptions.cmake index 852d522ce2..456f75a7ec 100644 --- a/cmake/FlagTreeOptions.cmake +++ b/cmake/FlagTreeOptions.cmake @@ -552,3 +552,14 @@ function(flagtree_spec_td_set output_td td_filename) endif() set(${output_td} ${ret} PARENT_SCOPE) endfunction() + + +# FLAGTREE SPEC MLIR DOC FUNC +function(flagtree_spec_add_mlir_doc + doc_filename output_file output_directory command) + flagtree_spec_td_set(_flagtree_doc_td "${doc_filename}.td") + string(REGEX REPLACE "\\.td$" "" _flagtree_doc_stem "${_flagtree_doc_td}") + add_mlir_doc( + "${_flagtree_doc_stem}" "${output_file}" "${output_directory}" + "${command}" ${ARGN}) +endfunction() diff --git a/include/triton/Dialect/Triton/IR/CMakeLists.txt b/include/triton/Dialect/Triton/IR/CMakeLists.txt index daece64408..63ff541af2 100644 --- a/include/triton/Dialect/Triton/IR/CMakeLists.txt +++ b/include/triton/Dialect/Triton/IR/CMakeLists.txt @@ -5,7 +5,7 @@ mlir_tablegen(Ops.h.inc -gen-op-decls) mlir_tablegen(Ops.cpp.inc -gen-op-defs) mlir_tablegen(OpsEnums.h.inc -gen-enum-decls) mlir_tablegen(OpsEnums.cpp.inc -gen-enum-defs) -add_mlir_doc(TritonOps TritonOps dialects/ -gen-op-doc) +flagtree_spec_add_mlir_doc(TritonOps TritonOps dialects/ -gen-op-doc) set(LLVM_TARGET_DEFINITIONS TritonDialect.td) mlir_tablegen(Dialect.h.inc -gen-dialect-decls) diff --git a/include/triton/Dialect/TritonGPU/IR/CMakeLists.txt b/include/triton/Dialect/TritonGPU/IR/CMakeLists.txt index c1d5946467..3d63f1271d 100644 --- a/include/triton/Dialect/TritonGPU/IR/CMakeLists.txt +++ b/include/triton/Dialect/TritonGPU/IR/CMakeLists.txt @@ -8,7 +8,7 @@ mlir_tablegen(Ops.cpp.inc -gen-op-defs) mlir_tablegen(Types.h.inc -gen-typedef-decls -typedefs-dialect=ttg) mlir_tablegen(Types.cpp.inc -gen-typedef-defs -typedefs-dialect=ttg) add_mlir_doc(TritonGPUDialect TritonGPUDialect dialects/ -gen-dialect-doc) -add_mlir_doc(TritonGPUOps TritonGPUOps dialects/ -gen-op-doc) +flagtree_spec_add_mlir_doc(TritonGPUOps TritonGPUOps dialects/ -gen-op-doc) add_public_tablegen_target(TritonGPUTableGen) flagtree_spec_td_set(LLVM_TARGET_DEFINITIONS TritonGPUAttrDefs.td) diff --git a/include/triton/Dialect/TritonInstrument/IR/CMakeLists.txt b/include/triton/Dialect/TritonInstrument/IR/CMakeLists.txt index fcf1486b00..827f25f839 100644 --- a/include/triton/Dialect/TritonInstrument/IR/CMakeLists.txt +++ b/include/triton/Dialect/TritonInstrument/IR/CMakeLists.txt @@ -10,6 +10,6 @@ mlir_tablegen(Ops.h.inc -gen-op-decls) mlir_tablegen(Ops.cpp.inc -gen-op-defs) mlir_tablegen(OpsEnums.h.inc -gen-enum-decls) mlir_tablegen(OpsEnums.cpp.inc -gen-enum-defs) -add_mlir_doc(TritonInstrumentOps TritonInstrumentOps dialects/ -gen-op-doc) +flagtree_spec_add_mlir_doc(TritonInstrumentOps TritonInstrumentOps dialects/ -gen-op-doc) add_public_tablegen_target(TritonInstrumentTableGen) diff --git a/include/triton/Dialect/TritonNvidiaGPU/IR/CMakeLists.txt b/include/triton/Dialect/TritonNvidiaGPU/IR/CMakeLists.txt index e36d751511..5c48006fa3 100644 --- a/include/triton/Dialect/TritonNvidiaGPU/IR/CMakeLists.txt +++ b/include/triton/Dialect/TritonNvidiaGPU/IR/CMakeLists.txt @@ -5,8 +5,8 @@ mlir_tablegen(Dialect.h.inc -gen-dialect-decls -dialect=ttng) mlir_tablegen(Dialect.cpp.inc -gen-dialect-defs -dialect=ttng) mlir_tablegen(Ops.h.inc -gen-op-decls) mlir_tablegen(Ops.cpp.inc -gen-op-defs) -add_mlir_doc(TritonNvidiaGPUDialect TritonNvidiaGPUDialect dialects/ -gen-dialect-doc) -add_mlir_doc(TritonNvidiaGPUOps TritonNvidiaGPUOps dialects/ -gen-op-doc) +flagtree_spec_add_mlir_doc(TritonNvidiaGPUDialect TritonNvidiaGPUDialect dialects/ -gen-dialect-doc) +flagtree_spec_add_mlir_doc(TritonNvidiaGPUOps TritonNvidiaGPUOps dialects/ -gen-op-doc) add_public_tablegen_target(TritonNvidiaGPUTableGen) flagtree_spec_td_set(LLVM_TARGET_DEFINITIONS TritonNvidiaGPUAttrDefs.td)