Skip to content
28 changes: 27 additions & 1 deletion mlir/lib/Conversion/ArithToSPIRV/ArithToSPIRV.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -607,6 +607,32 @@ struct UIToFPI1Pattern final : public OpConversionPattern<arith::UIToFPOp> {
}
};

//===----------------------------------------------------------------------===//
// IndexCastOp
//===----------------------------------------------------------------------===//

// Converts arith.index_cast to spirv.INotEqual if the target type is i1.
struct IndexCastIndexI1Pattern final
: public OpConversionPattern<arith::IndexCastOp> {
using OpConversionPattern::OpConversionPattern;

LogicalResult
matchAndRewrite(arith::IndexCastOp op, OpAdaptor adaptor,
ConversionPatternRewriter &rewriter) const override {
Type srcType = adaptor.getOperands().front().getType();
// Indexes have already been converted to its respective spirv type:
Type indexType = getTypeConverter<SPIRVTypeConverter>()->getIndexType();
if (srcType != indexType || !op.getType().isInteger(1))
return failure();

Location loc = op.getLoc();
Value zeroIdx = spirv::ConstantOp::getZero(srcType, loc, rewriter);
rewriter.replaceOpWithNewOp<spirv::INotEqualOp>(
op, rewriter.getI1Type(), zeroIdx, adaptor.getOperands().front());
return success();
}
};

//===----------------------------------------------------------------------===//
// ExtSIOp
//===----------------------------------------------------------------------===//
Expand Down Expand Up @@ -1328,7 +1354,7 @@ void mlir::arith::populateArithToSPIRVPatterns(
TypeCastingOpPattern<arith::SIToFPOp, spirv::ConvertSToFOp>,
TypeCastingOpPattern<arith::FPToUIOp, spirv::ConvertFToUOp>,
TypeCastingOpPattern<arith::FPToSIOp, spirv::ConvertFToSOp>,
TypeCastingOpPattern<arith::IndexCastOp, spirv::SConvertOp>,
TypeCastingOpPattern<arith::IndexCastOp, spirv::SConvertOp>, IndexCastIndexI1Pattern,
TypeCastingOpPattern<arith::IndexCastUIOp, spirv::UConvertOp>,
TypeCastingOpPattern<arith::BitcastOp, spirv::BitcastOp>,
CmpIOpBooleanPattern, CmpIOpPattern,
Expand Down
8 changes: 8 additions & 0 deletions mlir/test/Conversion/ArithToSPIRV/arith-to-spirv.mlir
Original file line number Diff line number Diff line change
Expand Up @@ -734,6 +734,14 @@ func.func @index_castui4(%arg0: index) {
return
}

// CHECK-LABEL: index_castindexi1
func.func @index_castindexi1(%arg0 : index) {
// CHECK: %[[ZERO:.+]] = spirv.Constant 0 : i32
// CHECK: spirv.INotEqual %[[ZERO]], %{{.+}} : i32
%0 = arith.index_cast %arg0 : index to i1
return
}

// CHECK-LABEL: @bit_cast
func.func @bit_cast(%arg0: vector<2xf32>, %arg1: i64) {
// CHECK: spirv.Bitcast %{{.+}} : vector<2xf32> to vector<2xi32>
Expand Down