Workaround

Author: nhat-nguyen

lib/AnalysisStructured/PtrAnalysis.cpp

@@ -231,8 +231,8 @@ LogicalResult PtrState::mulState(const PtrState &lhsState,
   }
 
   if (lhsState.scalar && rhsState.scalar) {
-    scalar = builder.create<arith::MulIOp>(
-        loc, lhsState.scalar, rhsState.scalar);
+    scalar =
+        builder.create<arith::MulIOp>(loc, lhsState.scalar, rhsState.scalar);
   }
 
   for (uint64_t i = 0; i < lhs->sizes.size(); i++) {
@@ -687,9 +687,15 @@ LogicalResult PtrAnalysis::visitOperand(Value operand, PtrState &state,
       } else if (auto makeTensorOp = dyn_cast<triton::MakeTensorPtrOp>(op)) {
         llvm_unreachable("Unexpected operand defining operation tts.make_tptr");
       } else {
+        op->dump();
+        return failure();
         llvm_unreachable("Unexpected operand defining operation");
       }
     } else {
+      OpBuilder::InsertionGuard guard(builder);
+      builder.setInsertionPointToStart(operand.getParentBlock());
+      state.scalar =
+          builder.create<arith::ConstantOp>(loc, builder.getIndexAttr(0));
       state.source = operand;
       return success();
     }

lib/Conversion/TritonToUnstructured/TritonToUnstructuredPass.cpp

@@ -269,6 +269,18 @@ class TritonToUnstructuredPass
       }
     });
 
+    getOperation().walk([&](triton::IntToPtrOp op) {
+      auto res = op.getResult();
+      OpBuilder b(op);
+      Value zero = b.create<arith::ConstantOp>(
+          op.getLoc(),
+          b.getIntegerAttr(IntegerType::get(&getContext(), defaultBitWidth),
+                           0));
+
+      offsetMap.insert({res, {res, res.getType(), defaultBitWidth, zero}});
+      workList.push(res);
+    });
+
     llvm::SmallVector<Operation *> toDelete;
     llvm::SmallVector<Operation *> ptrUsers;
 
@@ -350,6 +362,37 @@ class TritonToUnstructuredPass
 
                   return success();
                 })
+                .Case<triton::ExpandDimsOp>([&](ExpandDimsOp op) {
+                  auto res = op->getResult(0);
+                  auto resType = res.getType();
+
+                  if (!isPtrTypeLike(resType)) {
+                    return success();
+                  }
+
+                  auto ptr = op->getOperand(0);
+                  auto offsetInfo = offsetMap.at(ptr);
+
+                  OpBuilder b{op};
+                  auto clone =
+                      b.create(op->getLoc(), op->getName().getIdentifier(),
+                               ValueRange{offsetInfo.offset},
+                               TypeRange{getPtrOffsetType(
+                                   resType, offsetInfo.bitWidth)});
+
+                  PtrOffset newOffsetInfo{offsetInfo.ptr, resType,
+                                          offsetInfo.bitWidth,
+                                          clone->getResult(0)};
+
+                  offsetMap.insert({
+                      res,
+                      newOffsetInfo,
+                  });
+                  workList.push(res);
+                  toDelete.push_back(op);
+
+                  return success();
+                })
                 .Case<triton::LoadOp, triton::StoreOp, triton::MakeTensorPtrOp,
                       tts::MakeTensorPtrOp>([&](Operation *op) {
                   // Special case: