[mlir][spirv] Infer converted type of scf.for from the init value

[ThomasRaoux]

Summary

Instead of using the TypeConverter infer the value of the alloca created based on the init value. This will allow some ambiguous types like multidimensional vectors to be converted correctly.

[mravishankar]

mlir/lib/Conversion/SCFToSPIRV/SCFToSPIRV.cpp

188

Maybe add some context for this so that coming back to this later, we can know why it was done this way.

[ThomasRaoux]

Add extra comment

[ThomasRaoux]

Thanks Mahesh for the review. I added extra context explaining why we do it that way and we can revisit later.

mlir/lib/Conversion/SCFToSPIRV/SCFToSPIRV.cpp

188

Good point. Added more context.

mlir/lib/Conversion/SCFToSPIRV/SCFToSPIRV.cpp

 /// Helper function to replaces SCF op outputs with SPIR-V variable loads.
 /// We create VariableOp to handle the results value of the control flow region.
 /// spv.loop/spv.selection currently don't yield value. Right after the loop
 /// we load the value from the allocation and use it as the SCF op result.
 template <typename ScfOp, typename OpTy>
 static void replaceSCFOutputValue(ScfOp scfOp, OpTy newOp,
                                   SPIRVTypeConverter &typeConverter,
                                   ConversionPatternRewriter &rewriter,
-                                  ScfToSPIRVContextImpl *scfToSPIRVContext) {
+                                  ScfToSPIRVContextImpl *scfToSPIRVContext,
+                                  ArrayRef<Type> returnTypes) {
 
   Location loc = scfOp.getLoc();
   auto &allocas = scfToSPIRVContext->outputVars[newOp];
   SmallVector<Value, 8> resultValue;
-  for (Value result : scfOp.results()) {
-    auto convertedType = typeConverter.convertType(result.getType());
+  for (Type convertedType : returnTypes) {
     auto pointerType =
         spirv::PointerType::get(convertedType, spirv::StorageClass::Function);
     rewriter.setInsertionPoint(newOp);
     auto alloc = rewriter.create<spirv::VariableOp>(
         loc, pointerType, spirv::StorageClass::Function,
         /*initializer=*/nullptr);
     allocas.push_back(alloc);
     rewriter.setInsertionPointAfter(newOp);
ForOpConversion::matchAndRewrite(scf::ForOp forOp, ArrayRef<Value> operands,
   Block *continueBlock = loopOp.getContinueBlock();
   rewriter.setInsertionPointToEnd(continueBlock);
 
   // Add the step to the induction variable and branch to the header.
   Value updatedIndVar = rewriter.create<spirv::IAddOp>(
       loc, newIndVar.getType(), newIndVar, forOperands.step());
   rewriter.create<spirv::BranchOp>(loc, header, updatedIndVar);
 
+  // Infer the return types from the init operands. Vector type may get
+  // converted to CooperativeMatrix or to Vector type, to avoid having complex
+  // extra logic to figure out the right type we just infer it from the Init
+  // operands.
+  SmallVector<Type, 8> initTypes;
+  for (auto arg : forOperands.initArgs())
+    initTypes.push_back(arg.getType());
   replaceSCFOutputValue(forOp, loopOp, typeConverter, rewriter,
-                        scfToSPIRVContext);
+                        scfToSPIRVContext, initTypes);
   return success();
 }
 
 //===----------------------------------------------------------------------===//
 // scf::IfOp.
 //===----------------------------------------------------------------------===//
 
 LogicalResult
IfOpConversion::matchAndRewrite(scf::IfOp ifOp, ArrayRef<Value> operands,
   }
 
   // Create a `spv.BranchConditional` operation for selection header block.
   rewriter.setInsertionPointToEnd(selectionHeaderBlock);
   rewriter.create<spirv::BranchConditionalOp>(loc, ifOperands.condition(),
                                               thenBlock, ArrayRef<Value>(),
                                               elseBlock, ArrayRef<Value>());
 
+  SmallVector<Type, 8> returnTypes;
+  for (auto result : ifOp.results()) {
+    auto convertedType = typeConverter.convertType(result.getType());
+    returnTypes.push_back(convertedType);
+  }
   replaceSCFOutputValue(ifOp, selectionOp, typeConverter, rewriter,
-                        scfToSPIRVContext);
+                        scfToSPIRVContext, returnTypes);
   return success();
 }
 
 /// Yield is lowered to stores to the VariableOp created during lowering of the
 /// parent region. For loops we also need to update the branch looping back to
 /// the header with the loop carried values.
 LogicalResult TerminatorOpConversion::matchAndRewrite(
     scf::YieldOp terminatorOp, ArrayRef<Value> operands,