[mlir][StandardToSPIRV] Handle i1 case for lowering memref.load/store op

[hanchung]

Summary

This patch unconditionally converts i1 types to i8 types on memrefs. If the
extensions or capabilities are not met, they will be converted to i32. Hence the
logic in IntLoadPattern and IntStorePattern are also updated.

Also added the implementation of SPIRVTypeConverter::getOptions().

[hanchung]

fix debug build

[hanchung]

use build-in function

[hanchung]

I hit a really weird thing... and I can only repro with this PR...

module attributes {
  spv.target_env = #spv.target_env<
    #spv.vce<v1.0, [Int16, StorageBuffer16BitAccess, Shader],
    [SPV_KHR_storage_buffer_storage_class, SPV_KHR_16bit_storage]>, {}>
} {

func @load_i8(%arg0: memref<i8>) -> i8 {
  %0 = memref.load %arg0[] : memref<i8>
  return %0 : i8
}

}

The above one can be lowered to SPIR-V, but the below one can not

module attributes {
  spv.target_env = #spv.target_env<
    #spv.vce<v1.0, [Int16, StorageBuffer16BitAccess, Shader],
    [SPV_KHR_storage_buffer_storage_class, SPV_KHR_16bit_storage]>, {}>
} {

func @load_i1(%arg0: memref<i1>) -> i1 {
  %0 = memref.load %arg0[] : memref<i1>
  return %0 : i1
}

}

They are expected to generate the same IRs. And they do when there are no users.

The i1 case only works if it does not have users, like in the test.

I dumped operands[0] in ReturnOpPattern. i8 case showed the replaced op %15 = spv.ShiftRightArithmetic %14, %13 : i32, i32, and i1 case showed the old one %16 = memref.load <<UNKNOWN SSA VALUE>>[] : memref<i1>. It is weird because it said that the pattern applied successfully, but the users still get the old operand.

I also verified that i1 case and i8 case generate the same log with -debug and they are all marked legal and success.

I'd be appreciate if someone can help on this... It looks like a memref.load was created because the variable name is %16.

To repro, run mlir-opt -convert-std-to-spirv a.mlir

full log for i1: https://gist.github.com/hanhanW/8a4494f337629299959b6e7d59a882f5
full log for i8: https://gist.github.com/hanhanW/56a8fcdd6a8eb25a743a71667754e6e4

[hanchung]

@antiagainst is going to add some support to SPIR-V, and I will rebase on it and rework a bit.

[antiagainst]

I've created D100058 and D100059 to introduce controls over emulation. You can rebase after they are landed. With D100059 I think we can introduce another option there to control the # of bits for boolean types. So the getTypeNumBytes (which should probably be changed to getTypeNumBits) and all those convert*Type functions can properly handle i1.

[hanchung]

fix i1 issues

[hanchung]

Fixed, and enhanced the test that loads i1. (now it will return i1 to make sure there is no such issue again.)

[antiagainst]

Mostly okay to me. Major request to have an option for control.

mlir/lib/Conversion/StandardToSPIRV/StandardToSPIRV.cpp

995

s/isInt1/isBool/

1052

W already have typeConverter at L988.

mlir/lib/Dialect/SPIRV/Transforms/SPIRVConversion.cpp

344

s/convertMemrefI1Type/convertBoolMemrefType/

349

s/yte/yet/

377

Can you add an option in SPIRVTypeConverter::Options, like boolNumBits, to control this? (For now we can just support the case == 8 and bail out on type conversion for other cases.)

[hanchung]

address comments

mlir/lib/Conversion/StandardToSPIRV/StandardToSPIRV.cpp

1052

I did not notice it, thanks!

mlir/lib/Dialect/SPIRV/Transforms/SPIRVConversion.cpp

377

good point, it is useful

[hanchung]

Also added the implementation of SPIRVTypeConverter::getOptions().

[antiagainst]

Cool, just one final comment.

mlir/lib/Dialect/SPIRV/Transforms/SPIRVConversion.cpp

156

Oops, sorry I forgot about this. :) Thanks for fixing it!

365

I think here it should be

c++
if (numBoolBits != 8) {
  LLVM_DEBUG(llvm::dbgs() << "using non-8-bit storage for bool types unimplemented");
  return nullptr;
}

?

[hanchung]

address comments

mlir/lib/Dialect/SPIRV/Transforms/SPIRVConversion.cpp

365

good point! I feel that it will work, but did not test it. Let's have the check, and we can turn it back and add tests when we need it.

mlir/lib/Conversion/StandardToSPIRV/StandardToSPIRV.cpp

if (!memrefType.getElementType().isSignlessInteger())
     return failure();
 
   auto &typeConverter = *getTypeConverter<SPIRVTypeConverter>();
   spirv::AccessChainOp accessChainOp =
       spirv::getElementPtr(typeConverter, memrefType, loadOperands.memref(),
                            loadOperands.indices(), loc, rewriter);
 
   int srcBits = memrefType.getElementType().getIntOrFloatBitWidth();
+  // Treat memrefs on i1 as i8.
+  if (srcBits == 1)
+    srcBits = 8;
   auto dstType = typeConverter.convertType(memrefType)
                      .cast<spirv::PointerType>()
                      .getPointeeType()
                      .cast<spirv::StructType>()
                      .getElementType(0)
                      .cast<spirv::ArrayType>()
                      .getElementType();
   int dstBits = dstType.getIntOrFloatBitWidth();
IntLoadOpPattern::matchAndRewrite(memref::LoadOp loadOp,
   Value shiftValue =
       rewriter.create<spirv::ConstantOp>(loc, dstType, shiftValueAttr);
   result = rewriter.create<spirv::ShiftLeftLogicalOp>(loc, dstType, result,
                                                       shiftValue);
   result = rewriter.create<spirv::ShiftRightArithmeticOp>(loc, dstType, result,
                                                           shiftValue);
   rewriter.replaceOp(loadOp, result);
 
   assert(accessChainOp.use_empty());
   rewriter.eraseOp(accessChainOp);
 
   return success();
 }
 
 LogicalResult
 LoadOpPattern::matchAndRewrite(memref::LoadOp loadOp, ArrayRef<Value> operands,
                                ConversionPatternRewriter &rewriter) const {
if (!memrefType.getElementType().isSignlessInteger())
     return failure();
 
   auto loc = storeOp.getLoc();
   auto &typeConverter = *getTypeConverter<SPIRVTypeConverter>();
   spirv::AccessChainOp accessChainOp =
       spirv::getElementPtr(typeConverter, memrefType, storeOperands.memref(),
                            storeOperands.indices(), loc, rewriter);
   int srcBits = memrefType.getElementType().getIntOrFloatBitWidth();
+
+  // Treat memrefs on i1 as i8.
+  bool isInt1 = srcBits == 1;
+  if (isInt1)
+    srcBits = 8;
   auto dstType = typeConverter.convertType(memrefType)
                      .cast<spirv::PointerType>()
                      .getPointeeType()
                      .cast<spirv::StructType>()
                      .getElementType(0)
                      .cast<spirv::ArrayType>()
                      .getElementType();
   int dstBits = dstType.getIntOrFloatBitWidth();
IntStoreOpPattern::matchAndRewrite(memref::StoreOp storeOp,
   // Create a mask to clear the destination. E.g., if it is the second i8 in
   // i32, 0xFFFF00FF is created.
   Value mask = rewriter.create<spirv::ConstantOp>(
       loc, dstType, rewriter.getIntegerAttr(dstType, (1 << srcBits) - 1));
   Value clearBitsMask =
       rewriter.create<spirv::ShiftLeftLogicalOp>(loc, dstType, mask, offset);
   clearBitsMask = rewriter.create<spirv::NotOp>(loc, dstType, clearBitsMask);
 
-  Value storeVal =
-      shiftValue(loc, storeOperands.value(), offset, mask, dstBits, rewriter);
+  Value storeVal = storeOperands.value();
+  if (isInt1) {
+    Value zero = rewriter.create<ConstantIntOp>(loc, 0, dstType);
+    Value one = rewriter.create<ConstantIntOp>(loc, 1, dstType);
+    storeVal =
+        rewriter.create<spirv::SelectOp>(loc, dstType, storeVal, one, zero);
+  }
+  storeVal = shiftValue(loc, storeVal, offset, mask, dstBits, rewriter);
   Value adjustedPtr = adjustAccessChainForBitwidth(typeConverter, accessChainOp,
                                                    srcBits, dstBits, rewriter);
   Optional<spirv::Scope> scope = getAtomicOpScope(memrefType);
   if (!scope)
     return failure();
   Value result = rewriter.create<spirv::AtomicAndOp>(
       loc, dstType, adjustedPtr, *scope, spirv::MemorySemantics::AcquireRelease,
       clearBitsMask);

mlir/lib/Dialect/SPIRV/Transforms/SPIRVConversion.cpp

#define STORAGE_SPACE_MAP_FN(storage, space)                                   \
   switch (space) {
     STORAGE_SPACE_MAP_LIST(STORAGE_SPACE_MAP_FN)
   default:
     return llvm::None;
   }
 #undef STORAGE_SPACE_MAP_FN
 }
 
 #undef STORAGE_SPACE_MAP_LIST
 
 // TODO: This is a utility function that should probably be
 // exposed by the SPIR-V dialect. Keeping it local till the use case arises.
 static Optional<int64_t> getTypeNumBytes(Type t) {
   if (t.isa<spirv::ScalarType>()) {
     auto bitWidth = t.getIntOrFloatBitWidth();
     // According to the SPIR-V spec:
     // "There is no physical size or bit pattern defined for values with boolean
if (!arrayElemSize) {
     LLVM_DEBUG(llvm::dbgs()
                << type << " illegal: cannot deduce converted element size\n");
     return llvm::None;
   }
 
   return spirv::ArrayType::get(*arrayElemType, arrayElemCount, *arrayElemSize);
 }
 
 static Optional<Type> convertMemrefType(const spirv::TargetEnv &targetEnv,
                                         MemRefType type) {
   Optional<spirv::StorageClass> storageClass =
       SPIRVTypeConverter::getStorageClassForMemorySpace(
           type.getMemorySpaceAsInt());
   if (!storageClass) {
     LLVM_DEBUG(llvm::dbgs()
                << type << " illegal: cannot convert memory space\n");
     return llvm::None;
   }
 
   Optional<Type> arrayElemType;
   Type elementType = type.getElementType();
   if (auto vecType = elementType.dyn_cast<VectorType>()) {
     arrayElemType = convertVectorType(targetEnv, vecType, storageClass);
   } else if (auto scalarType = elementType.dyn_cast<spirv::ScalarType>()) {
     arrayElemType = convertScalarType(targetEnv, scalarType, storageClass);
   } else {
     LLVM_DEBUG(
         llvm::dbgs()
         << type
         << " unhandled: can only convert scalar or vector element type\n");
     return llvm::None;
   }
   if (!arrayElemType)
     return llvm::None;
 
   Optional<int64_t> elementSize = getTypeNumBytes(elementType);
   if (!elementSize) {
     LLVM_DEBUG(llvm::dbgs()
                << type << " illegal: cannot deduce element size\n");
     return llvm::None;
   }
 
   if (!type.hasStaticShape()) {
     auto arrayType = spirv::RuntimeArrayType::get(*arrayElemType, *elementSize);
     // Wrap in a struct to satisfy Vulkan interface requirements.
     auto structType = spirv::StructType::get(arrayType, 0);
     return spirv::PointerType::get(structType, *storageClass);
   }
 
   Optional<int64_t> memrefSize = getTypeNumBytes(type);
addConversion([this](VectorType vectorType) {
     return convertVectorType(targetEnv, vectorType);
   });
 
   addConversion([this](TensorType tensorType) {
     return convertTensorType(targetEnv, tensorType);
   });
 
   addConversion([this](MemRefType memRefType) {
+    // Treat i1 on memrefs as i8.
+    if (memRefType.getElementTypeBitWidth() == 1) {
+      return convertMemrefType(
+          targetEnv,
+          MemRefType::get(memRefType.getShape(),
+                          IntegerType::get(memRefType.getContext(), 8)));
+    }
     return convertMemrefType(targetEnv, memRefType);
   });
 }
 
 //===----------------------------------------------------------------------===//
 // FuncOp Conversion Patterns
 //===----------------------------------------------------------------------===//
 

mlir/test/Conversion/StandardToSPIRV/std-ops-to-spirv.mlir

 }
 
 } // end module
 
 // -----
 
 // Check that access chain indices are properly adjusted if non-32-bit types are
 // emulated via 32-bit types.
-// TODO: Test i1 and i64 types.
+// TODO: Test i64 types.
 module attributes {
   spv.target_env = #spv.target_env<
     #spv.vce<v1.0, [Shader], [SPV_KHR_storage_buffer_storage_class]>, {}>
 } {
 
+// CHECK-LABEL: @load_i1
+func @load_i1(%arg0: memref<i1>) {
+  //     CHECK: %[[ZERO:.+]] = spv.Constant 0 : i32
+  //     CHECK: %[[FOUR1:.+]] = spv.Constant 4 : i32
+  //     CHECK: %[[QUOTIENT:.+]] = spv.SDiv %[[ZERO]], %[[FOUR1]] : i32
+  //     CHECK: %[[PTR:.+]] = spv.AccessChain %{{.+}}[%[[ZERO]], %[[QUOTIENT]]]
+  //     CHECK: %[[LOAD:.+]] = spv.Load  "StorageBuffer" %[[PTR]]
+  //     CHECK: %[[FOUR2:.+]] = spv.Constant 4 : i32
+  //     CHECK: %[[EIGHT:.+]] = spv.Constant 8 : i32
+  //     CHECK: %[[IDX:.+]] = spv.UMod %[[ZERO]], %[[FOUR2]] : i32
+  //     CHECK: %[[BITS:.+]] = spv.IMul %[[IDX]], %[[EIGHT]] : i32
+  //     CHECK: %[[VALUE:.+]] = spv.ShiftRightArithmetic %[[LOAD]], %[[BITS]] : i32, i32
+  //     CHECK: %[[MASK:.+]] = spv.Constant 255 : i32
+  //     CHECK: %[[T1:.+]] = spv.BitwiseAnd %[[VALUE]], %[[MASK]] : i32
+  //     CHECK: %[[T2:.+]] = spv.Constant 24 : i32
+  //     CHECK: %[[T3:.+]] = spv.ShiftLeftLogical %[[T1]], %[[T2]] : i32, i32
+  //     CHECK: spv.ShiftRightArithmetic %[[T3]], %[[T2]] : i32, i32
+  %0 = memref.load %arg0[] : memref<i1>
+  return
+}
+
 // CHECK-LABEL: @load_i8
 func @load_i8(%arg0: memref<i8>) {
   //     CHECK: %[[ZERO:.+]] = spv.Constant 0 : i32
   //     CHECK: %[[FOUR1:.+]] = spv.Constant 4 : i32
   //     CHECK: %[[QUOTIENT:.+]] = spv.SDiv %[[ZERO]], %[[FOUR1]] : i32
   //     CHECK: %[[PTR:.+]] = spv.AccessChain %{{.+}}[%[[ZERO]], %[[QUOTIENT]]]
   //     CHECK: %[[LOAD:.+]] = spv.Load  "StorageBuffer" %[[PTR]]
   //     CHECK: %[[FOUR2:.+]] = spv.Constant 4 : i32

 func @load_f32(%arg0: memref<f32>) {
   // CHECK-NOT: spv.SDiv
   //     CHECK: spv.Load
   // CHECK-NOT: spv.ShiftRightArithmetic
   %0 = memref.load %arg0[] : memref<f32>
   return
 }
 
+// CHECK-LABEL: @store_i1
+//       CHECK: (%[[ARG0:.+]]: {{.*}}, %[[ARG1:.+]]: i1)
+func @store_i1(%arg0: memref<i1>, %value: i1) {
+  //     CHECK: %[[ZERO:.+]] = spv.Constant 0 : i32
+  //     CHECK: %[[FOUR:.+]] = spv.Constant 4 : i32
+  //     CHECK: %[[EIGHT:.+]] = spv.Constant 8 : i32
+  //     CHECK: %[[IDX:.+]] = spv.UMod %[[ZERO]], %[[FOUR]] : i32
+  //     CHECK: %[[OFFSET:.+]] = spv.IMul %[[IDX]], %[[EIGHT]] : i32
+  //     CHECK: %[[MASK1:.+]] = spv.Constant 255 : i32
+  //     CHECK: %[[TMP1:.+]] = spv.ShiftLeftLogical %[[MASK1]], %[[OFFSET]] : i32, i32
+  //     CHECK: %[[MASK:.+]] = spv.Not %[[TMP1]] : i32
+  //     CHECK: %[[ZERO1:.+]] = spv.Constant 0 : i32
+  //     CHECK: %[[ONE1:.+]] = spv.Constant 1 : i32
+  //     CHECK: %[[CASTED_ARG1:.+]] = spv.Select %[[ARG1]], %[[ONE1]], %[[ZERO1]] : i1, i32
+  //     CHECK: %[[CLAMPED_VAL:.+]] = spv.BitwiseAnd %[[CASTED_ARG1]], %[[MASK1]] : i32
+  //     CHECK: %[[STORE_VAL:.+]] = spv.ShiftLeftLogical %[[CLAMPED_VAL]], %[[OFFSET]] : i32, i32
+  //     CHECK: %[[FOUR2:.+]] = spv.Constant 4 : i32
+  //     CHECK: %[[ACCESS_IDX:.+]] = spv.SDiv %[[ZERO]], %[[FOUR2]] : i32
+  //     CHECK: %[[PTR:.+]] = spv.AccessChain %[[ARG0]][%[[ZERO]], %[[ACCESS_IDX]]]
+  //     CHECK: spv.AtomicAnd "Device" "AcquireRelease" %[[PTR]], %[[MASK]]
+  //     CHECK: spv.AtomicOr "Device" "AcquireRelease" %[[PTR]], %[[STORE_VAL]]
+  memref.store %value, %arg0[] : memref<i1>
+  return
+}
+
 // CHECK-LABEL: @store_i8
 //       CHECK: (%[[ARG0:.+]]: {{.*}}, %[[ARG1:.+]]: i32)
 func @store_i8(%arg0: memref<i8>, %value: i8) {
   //     CHECK: %[[ZERO:.+]] = spv.Constant 0 : i32
   //     CHECK: %[[FOUR:.+]] = spv.Constant 4 : i32
   //     CHECK: %[[EIGHT:.+]] = spv.Constant 8 : i32
   //     CHECK: %[[IDX:.+]] = spv.UMod %[[ZERO]], %[[FOUR]] : i32
   //     CHECK: %[[OFFSET:.+]] = spv.IMul %[[IDX]], %[[EIGHT]] : i32

mlir/test/Conversion/StandardToSPIRV/std-types-to-spirv.mlir

func @memref_mem_space(
     %arg4: memref<4xf32, 5>,
     %arg5: memref<4xf32, 6>
 ) { return }
 
 } // end module
 
 // -----
 
-// Check that boolean memref is not supported at the moment.
-module attributes {
-  spv.target_env = #spv.target_env<#spv.vce<v1.0, [], []>, {}>
-} {
-
-// CHECK-LABEL: func @memref_type({{%.*}}: memref<3xi1>)
-func @memref_type(%arg0: memref<3xi1>) {
-  return
-}
-
-} // end module
-
-// -----
-
 // Check that using non-32-bit scalar types in interface storage classes
 // requires special capability and extension: convert them to 32-bit if not
 // satisfied.
 module attributes {
   spv.target_env = #spv.target_env<#spv.vce<v1.0, [], []>, {}>
 } {
 
+// CHECK-LABEL: spv.func @memref_1bit_type
+// CHECK-SAME: !spv.ptr<!spv.struct<(!spv.array<3 x i32, stride=4> [0])>, StorageBuffer>
+func @memref_1bit_type(%arg0: memref<3xi1>) { return }
+
 // CHECK-LABEL: spv.func @memref_8bit_StorageBuffer
 // CHECK-SAME: !spv.ptr<!spv.struct<(!spv.array<16 x i32, stride=4> [0])>, StorageBuffer>
 func @memref_8bit_StorageBuffer(%arg0: memref<16xi8, 0>) { return }
 
 // CHECK-LABEL: spv.func @memref_8bit_Uniform
 // CHECK-SAME: !spv.ptr<!spv.struct<(!spv.array<16 x si32, stride=4> [0])>, Uniform>
 func @memref_8bit_Uniform(%arg0: memref<16xsi8, 4>) { return }