Skip to content

Commit c6d7faa

Browse files
pthomadakispthomadakis
authored
[GPU] Support explicit block sizes, SPMM (#109)
* Support explicitly setting block_size and parallel dims in GPU * Reenabed support for SPMM * Added attribute in IndexTreeIndicesOp that is used to map an iteration space to a GPU block * Modified ParallelLoopsToGpu to operation on scf::Forall instead of scf::Parallel * Changed mlir generator to automatically generate i32 indices for sparse matrices when targetting GPU devices * Updated cometPy to latestt GPU changes
1 parent 47963b1 commit c6d7faa

25 files changed

Lines changed: 1257 additions & 207 deletions

File tree

‎frontends/comet_dsl/CMakeLists.txt‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -53,7 +53,7 @@ if(ENABLE_GPU_TARGET)
5353
set(LIBS
5454
${LIBS}
5555
COMETGPUUtils
56-
COMETParallelLoopsToGpu
56+
COMETForallToGpu
5757
COMETGpuToBlockedGpu
5858
COMETBlockedGpuToTriton
5959
#COMETGpuToTriton

‎frontends/comet_dsl/comet.cpp‎

Lines changed: 28 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -106,7 +106,7 @@
106106
#include "comet/Conversion/PrepareGpuHost/PrepareGpuHostPass.h"
107107
#include "comet/Conversion/BlockedGpuToTriton/BlockedGpuToTriton.h"
108108
#include "comet/Conversion/GpuToBlockedGpu/GpuToBlockedGpu.h"
109-
#include "comet/Conversion/ParallelLoopsToGpu/ParallelLoopsToGpu.h"
109+
#include "comet/Conversion/ForallToGpu/ForallToGpu.h"
110110
#include "comet/Conversion/TritonToHIP/TritonToHIPPass.h"
111111
#include "triton/Dialect/TritonGPU/IR/Dialect.h"
112112
#include "mlir/Conversion/SCFToGPU/SCFToGPUPass.h"
@@ -340,7 +340,7 @@ std::unique_ptr<tensorAlgebra::ModuleAST> parseInputFile(llvm::StringRef filenam
340340
}
341341

342342
int loadMLIR(mlir::MLIRContext &context,
343-
mlir::OwningOpRef<mlir::ModuleOp> &module)
343+
mlir::OwningOpRef<mlir::ModuleOp> &module, bool useI64)
344344
{
345345
/// Handle '.ta' input to the compiler.
346346
if (inputType != InputType::MLIR &&
@@ -349,7 +349,7 @@ int loadMLIR(mlir::MLIRContext &context,
349349
auto moduleAST = parseInputFile(inputFilename);
350350
if (!moduleAST)
351351
return 6;
352-
module = mlirGen(context, *moduleAST);
352+
module = mlirGen(context, *moduleAST, useI64);
353353
return !module ? 1 : 0;
354354
}
355355

@@ -375,7 +375,7 @@ int loadMLIR(mlir::MLIRContext &context,
375375
}
376376

377377
int loadAndProcessMLIR(mlir::MLIRContext &context,
378-
mlir::OwningOpRef<mlir::ModuleOp> &module)
378+
mlir::OwningOpRef<mlir::ModuleOp> &module, bool useI64)
379379
{
380380
#ifdef ENABLE_GPU_TARGET
381381
bool emitTriton_ = emitTriton && CodegenTarget == TargetDevice::GPU;
@@ -389,7 +389,7 @@ int loadAndProcessMLIR(mlir::MLIRContext &context,
389389
tensorAlgebra::debugOptions.insert("debug-ta-labels-alphabet-order");
390390
}
391391
/// end Load debug options
392-
if (int error = loadMLIR(context, module))
392+
if (int error = loadMLIR(context, module, useI64))
393393
return error;
394394

395395
mlir::PassManager pm(module.get()->getName());
@@ -596,9 +596,19 @@ int loadAndProcessMLIR(mlir::MLIRContext &context,
596596
/// Blanket-convert any remaining affine ops if any remain.
597597
pm.addPass(mlir::createLowerAffinePass());
598598
/// Convert SCF to CF (always needed).
599-
pm.addPass(mlir::createForallToParallelLoopPass());
600-
pm.addPass(mlir::createLoopInvariantCodeMotionPass());
601-
pm.addPass(mlir::createCanonicalizerPass());
599+
#ifdef ENABLE_GPU_TARGET
600+
if(CodegenTarget != TargetDevice::GPU)
601+
{
602+
pm.addPass(mlir::createForallToParallelLoopPass());
603+
pm.addPass(mlir::createLoopInvariantCodeMotionPass());
604+
pm.addPass(mlir::createCanonicalizerPass());
605+
}
606+
#else
607+
pm.addPass(mlir::createForallToParallelLoopPass());
608+
pm.addPass(mlir::createLoopInvariantCodeMotionPass());
609+
pm.addPass(mlir::createCanonicalizerPass());
610+
#endif
611+
602612

603613
#ifndef ENABLE_GPU_TARGET
604614
[[maybe_unused]] bool IsLoweringToTriton = false;
@@ -615,7 +625,7 @@ int loadAndProcessMLIR(mlir::MLIRContext &context,
615625
#ifdef ENABLE_GPU_TARGET
616626
if (CodegenTarget == TargetDevice::GPU)
617627
{
618-
pm.addNestedPass<mlir::func::FuncOp>(mlir::comet::createConvertParallelLoopsToGpuPass(GPUBlockSizeX, GPUBlockSizeY, GPUBlockSizeR));
628+
pm.addNestedPass<mlir::func::FuncOp>(mlir::comet::createConvertForallToGpuPass(GPUBlockSizeX, GPUBlockSizeY, GPUBlockSizeR));
619629
}
620630
#endif
621631

@@ -804,8 +814,16 @@ int main(int argc, char **argv)
804814
context.loadDialect<mlir::index::IndexDialect>();
805815

806816
mlir::OwningOpRef<mlir::ModuleOp> module;
817+
bool useI64 = true;
818+
#ifdef ENABLE_GPU_TARGET
819+
if(CodegenTarget == TargetDevice::GPU)
820+
{
821+
useI64 = false;
822+
}
823+
#endif
824+
807825

808-
if (int error = loadAndProcessMLIR(context, module))
826+
if (int error = loadAndProcessMLIR(context, module, useI64))
809827
return error;
810828

811829
/// If we aren't exporting to non-mlir, then we are done.

‎frontends/comet_dsl/include/MLIRGen.h‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -48,7 +48,7 @@ namespace tensorAlgebra
4848
/// Emit IR for the given Tensor Algebra moduleAST, returns a newly created MLIR module
4949
/// or nullptr on failure.
5050
mlir::OwningOpRef<mlir::ModuleOp> mlirGen(mlir::MLIRContext &context,
51-
ModuleAST &moduleAST);
51+
ModuleAST &moduleAST, bool useI64);
5252

5353
extern std::unordered_set<std::string> debugOptions;
5454
} /// namespace tensorAlgebra

‎frontends/comet_dsl/mlir/MLIRGen.cpp‎

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2515,8 +2515,12 @@ namespace tensorAlgebra
25152515

25162516
/// The public API for codegen.
25172517
mlir::OwningOpRef<mlir::ModuleOp> mlirGen(mlir::MLIRContext &context,
2518-
ModuleAST &moduleAST)
2518+
ModuleAST &moduleAST, bool useI64)
25192519
{
2520+
if(!useI64)
2521+
{
2522+
defaultSpTensorIndiceBitWidth = 32;
2523+
}
25202524
return MLIRGenImpl(context).mlirGen(moduleAST);
25212525
}
25222526

‎frontends/numpy-scipy/cometpy/MLIRGen/ops.py‎

Lines changed: 19 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -468,6 +468,24 @@ def dump(self):
468468
undefined=jinja2.StrictUndefined,
469469
)
470470

471+
472+
class ForallInParallel(Operation):
473+
474+
def __init__(self, values):
475+
super().__init__(values, [v.type for v in values], endsBody=True)
476+
self. values = values
477+
478+
def dump(self):
479+
return self.text.render(
480+
values = ",".join([f'%{v}' for v in self.values]),
481+
)
482+
483+
484+
text = jinja2.Template(
485+
'scf.forall.in_parallel {}',
486+
undefined=jinja2.StrictUndefined,
487+
)
488+
471489

472490
class ReturnOp(Operation):
473491

@@ -603,6 +621,6 @@ def dump(self):
603621

604622

605623
text = jinja2.Template(
606-
'scf.parallel (%{{iv}}) = (%{{lb}}) to (%{{ub}}) step (%{{step}}) ',
624+
'scf.forall (%{{iv}}) = (%{{lb}}) to (%{{ub}}) step (%{{step}}) ',
607625
undefined=jinja2.StrictUndefined,
608626
)

‎frontends/numpy-scipy/cometpy/comet.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -413,7 +413,7 @@ def visit_For(self, node):
413413
for stmt in node.body:
414414
self.visit(stmt)
415415
if pragma == "#pragma parallel":
416-
self.build(ops.ReduceOp([]))
416+
self.build(ops.ForallInParallel([]))
417417
else:
418418
self.build(ops.YieldOp([]))
419419

‎include/comet/Conversion/CMakeLists.txt‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
if(ENABLE_GPU_TARGET OR ENABLE_FPGA_TARGET)
2-
add_subdirectory(ParallelLoopsToGpu)
2+
add_subdirectory(ForallToGpu)
33
add_subdirectory(ParallelLoopsToGpuFPGA)
44
endif()
55
if(ENABLE_FPGA_TARGET)

‎include/comet/Conversion/ForallToGpu/CMakeLists.txt‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,3 @@
1+
set(LLVM_TARGET_DEFINITIONS Passes.td)
2+
mlir_tablegen(Passes.h.inc -gen-pass-decls --name ForallToGpu)
3+
add_public_tablegen_target(ForallConversionPassIncGen)

include/comet/Conversion/ParallelLoopsToGpu/ParallelLoopsToGpu.h renamed to include/comet/Conversion/ForallToGpu/ForallToGpu.h

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -19,8 +19,8 @@
1919
// OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
2020
//
2121

22-
#ifndef COMET_CONVERSION_PARALLELLOOPSTOGPU_H
23-
#define COMET_CONVERSION_PARALLELLOOPSTOGPU_H
22+
#ifndef COMET_CONVERSION_FORALLTOGPU_H
23+
#define COMET_CONVERSION_FORALLTOGPU_H
2424

2525
#include <memory>
2626
namespace mlir {
@@ -30,8 +30,8 @@ class FuncOp;
3030
}
3131
template <typename T> class OperationPass;
3232
namespace comet {
33-
std::unique_ptr<OperationPass<mlir::func::FuncOp>> createConvertParallelLoopsToGpuPass();
34-
std::unique_ptr<OperationPass<mlir::func::FuncOp>> createConvertParallelLoopsToGpuPass(int blockX, int blockY, int blockR);
33+
std::unique_ptr<OperationPass<mlir::func::FuncOp>> createConvertForallToGpuPass();
34+
std::unique_ptr<OperationPass<mlir::func::FuncOp>> createConvertForallToGpuPass(int blockX, int blockY, int blockR);
3535
}
3636
}
3737

include/comet/Conversion/ParallelLoopsToGpu/Passes.h renamed to include/comet/Conversion/ForallToGpu/Passes.h

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -19,15 +19,15 @@
1919
// OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
2020
//
2121

22-
#ifndef COMET_PARALLEL_LOOPS_TO_GPU_CONVERSION_PASSES
23-
#define COMET_PARALLEL_LOOPS_TO_GPU_CONVERSION_PASSES
22+
#ifndef COMET_FORALL_TO_GPU_CONVERSION_PASSES
23+
#define COMET_FORALL_TO_GPU_CONVERSION_PASSES
2424

25-
#include "comet/Conversion/ParallelLoopsToGpu/ParallelLoopsToGpu.h"
25+
#include "comet/Conversion/ForallToGpu/ForallToGpu.h"
2626

2727
namespace mlir {
2828
namespace comet {
2929
#define GEN_PASS_REGISTRATION
30-
#include "comet/Conversion/ParallelLoopsToGpu/Passes.h.inc"
30+
#include "comet/Conversion/ForallToGpu/Passes.h.inc"
3131
}
3232
}
3333

0 commit comments

Comments
 (0)