Fix build failures from cfe3dd0

Signed-off-by: Whitney Tsang <[email protected]>

Author: whitneywhtsang

third_party/intel/lib/TritonIntelGPUToLLVM/PrintOpToLLVM.cpp

@@ -29,11 +29,10 @@ struct PrintOpConversion
                   ConversionPatternRewriter &rewriter) const override {
     auto loc = op->getLoc();
 
-    auto getPid = [&](int axis) {
-      return targetInfo.programId(rewriter, loc,
-                                  op->getParentOfType<ModuleOp>(), axis);
-    };
-    std::array<Value, 3> pid = {getPid(0), getPid(1), getPid(2)};
+    std::array<Value, 3> pid;
+    auto module = op->getParentOfType<ModuleOp>();
+    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) {

third_party/intel/lib/TritonIntelGPUToLLVM/TargetInfo.cpp

@@ -86,16 +86,9 @@ Value TargetInfo::shuffleIdx(RewriterBase &rewriter, Location loc, Value val,
 }
 
 Value TargetInfo::programId(RewriterBase &rewriter, Location loc,
-                            ModuleOp moduleOp, int axis) const {
-  assert(axis >= 0);
-  assert(axis < 3);
-  assert(moduleOp);
-
-  constexpr mlir::gpu::Dimension dims[] = {mlir::gpu::Dimension::x,
-                                           mlir::gpu::Dimension::y,
-                                           mlir::gpu::Dimension::z};
-
-  Value blockId = rewriter.create<::mlir::gpu::BlockIdOp>(loc, dims[axis]);
+                            ModuleOp moduleOp, ProgramIDDim axis) const {
+  Value blockId =
+      rewriter.create<::mlir::gpu::BlockIdOp>(loc, mlir::gpu::Dimension(axis));
   return rewriter.create<arith::IndexCastOp>(loc, i32_ty, blockId);
 }
 

third_party/intel/lib/TritonIntelGPUToLLVM/TargetInfo.h

@@ -42,7 +42,7 @@ class TargetInfo : public mlir::triton::TargetInfoBase {
                    Value i) const override;
 
   Value programId(RewriterBase &rewriter, Location loc, ModuleOp moduleOp,
-                  int axis) const override;
+                  ProgramIDDim axis) const override;
 
   bool warpReduce(RewriterBase &rewriter, Location loc, SmallVector<Value> &acc,
                   triton::ReduceOp op, unsigned numLaneToReduce,