Lines Matching defs:spmdizationMap

179     ArrayRef<MeshAxis> reductionMeshAxes, IRMapping &spmdizationMap,
187       spmdizationMap.lookup(op->getOperands()[operandIdx]);
196     IRMapping &spmdizationMap, ImplicitLocOpBuilder &builder) {
207   Value spmdizedLinalgOpResult = spmdizationMap.lookup(unshardedLinalgOpResult);
211   spmdizationMap.map(unshardedLinalgOpResult, reducedValue);
216     ArrayRef<MeshSharding> resultShardings, IRMapping &spmdizationMap,
223         reductionKind, spmdizationMap, builder);
233     IRMapping &spmdizationMap, SymbolTableCollection &symbolTable,
241                                                 spmdizationMap, builder);
242   // We must not change the operand mappings of the original spmdizationMap as
254     spmdizationMap.map(result, internalSpmdizationMap.lookup(result));
259       op, reductionMeshAxes, resultShardings, spmdizationMap, builder);
305                         IRMapping &spmdizationMap,
331           loopIteratorTypes, meshAxisAssignmentForLoopIterators, spmdizationMap,
336                                          spmdizationMap, symbolTable, builder);