Skip to content

[MLIR][XeGPU][TransformOps] Add slice_dims argument to set_op_layout_attr and set_desc_layout #168929

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom
Merged
merged 1 commit into from
Nov 21, 2025
Merged

[MLIR][XeGPU][TransformOps] Add slice_dims argument to set_op_layout_attr and set_desc_layout #168929

merged 1 commit into from
Nov 21, 2025
+139 −15

Conversation

@tkarna
Copy link
Contributor

@tkarna tkarna commented Nov 20, 2025 •
edited
Loading

set_op_layout_attr and set_desc_layout transform ops wrap xegpu.layout in an xegpu.slice attribute if slice_dims argument is set. This is needed for annotating vector.broadcast operations, for example. slice_dims only accepts integers because tensor ranks are usually fixed and thus there's no need to parametrize the dimensions.

@llvmbot llvmbot added mlir:gpu mlir:python MLIR Python bindings mlir labels Nov 20, 2025
@llvmbot
Copy link
Member

llvmbot commented Nov 20, 2025

@llvm/pr-subscribers-mlir-gpu

Author: Tuomas Kärnä (tkarna)

Changes

set_op_layout_attr and set_desc_layout transform ops wrap xegpu.layout in an xegpu.slice attribute if slice_dims argument is set. This is needed for annotating vector.broadcast operations, for example.

Patch is 53.61 KiB, truncated to 20.00 KiB below, full version: https://github.com/llvm/llvm-project/pull/168929.diff

5 Files Affected:

  • (modified) mlir/include/mlir/Dialect/XeGPU/TransformOps/XeGPUTransformOps.td (+17-7)
  • (modified) mlir/lib/Dialect/XeGPU/TransformOps/XeGPUTransformOps.cpp (+28-7)
  • (modified) mlir/python/mlir/dialects/transform/xegpu.py (+8)
  • (modified) mlir/test/Dialect/XeGPU/transform-ops.mlir (+38)
  • (modified) mlir/test/python/dialects/transform_xegpu_ext.py (+48-1)
<!DOCTYPE html>
<!--

Hello future GitHubber! I bet you're here to remove those nasty inline styles,
DRY up these templates and make 'em nice and re-usable, right?

Please, don't. https://github.com/styleguide/templates/2.0

-->
<html>
  <head>
    <title>Unicorn! &middot; GitHub</title>
    <style type="text/css" media="screen">
      body {
        background-color: #f1f1f1;
        margin: 0;
        font-family: "Helvetica Neue", Helvetica, Arial, sans-serif;
      }

      .container { margin: 50px auto 40px auto; width: 600px; text-align: center; }

      a { color: #4183c4; text-decoration: none; }
      a:hover { text-decoration: underline; }

      h1 { letter-spacing: -1px; line-height: 60px; font-size: 60px; font-weight: 100; margin: 0px; text-shadow: 0 1px 0 #fff; }
      p { color: rgba(0, 0, 0, 0.5); margin: 10px 0 10px; font-size: 18px; font-weight: 200; line-height: 1.6em;}

      ul { list-style: none; margin: 25px 0; padding: 0; }
      li { display: table-cell; font-weight: bold; width: 1%; }

      .logo { display: inline-block; margin-top: 35px; }
      .logo-img-2x { display: none; }
      @media
      only screen and (-webkit-min-device-pixel-ratio: 2),
      only screen and (   min--moz-device-pixel-ratio: 2),
      only screen and (     -o-min-device-pixel-ratio: 2/1),
      only screen and (        min-device-pixel-ratio: 2),
      only screen and (                min-resolution: 192dpi),
      only screen and (                min-resolution: 2dppx) {
        .logo-img-1x { display: none; }
        .logo-img-2x { display: inline-block; }
      }

      #suggestions {
        margin-top: 35px;
        color: #ccc;
      }
      #suggestions a {
        color: #666666;
        font-weight: 200;
        font-size: 14px;
        margin: 0 10px;
      }

    </style>
  </head>
  <body>

    <div class="container">
      <p>
        <img width="200" src="data:image/png;base64,
iVBORw0KGgoAAAANSUhEUgAAAZAAAAGZCAMAAACQbpc2AAADAFBMVEWEBz6FAD6FAD6GAD+MAEGOAEOKAEGOAEOGAD+IAECOAEOOAUOOAEOOAEOOAEOOAEOOAEOOAEOCACyOAEQAAACKAEJpoJ2KADqu0eSKAD2BAD+KAD6AxCNqwoX1Ziawo9LYnGuwhL7aiwaIyYudst6DlrnhcWvCcVvlbE9PvsL2u3uYmdCZum9Rns7clUYsreOYp2nuj37kiZLUfqhvxafqfmrNWEjA22uz1Vu6lsiRAArWhLfgdoONnF36rWrzpz/ajy/LbqR8pdWDVXx8faN6m0+6hE3+1pmHP2mTy1b4k1p/ueFlj7vMfzyBs1v/2GfLpFL5rQD+7teQABj4n0/rpCz6r1Z8apP+4rrfw1iV0JyLUkqQACTDnah5AACtaEt9zcGNhleJbU5lsOFRvE5FuOPS4WLSt7+8ajvXcJpcyNvhztSwUkTGgZz9uDSxRnPZwsnTkqy7jpsPt/AxyvnAeJMAxPnMiqPJYIyzf4/KqrPedaK5bIsQwfTo2t72hUepPWvn6mO8VoXrgbDRaJXle6r75FzZm7S2TnyTNUXv5eiDAA6j0WmKyWv51liRAACjP0OZNVuwa4L+uQAcw/SiNGH57mIxxPKmU3DgpL2sX3rutc3/vh9gu1D3tNEAtvOSKVAmxfagRmbnrcWZJVaTJET2qsv84eyk04PzcEb71OX0u9OUzHD1nsR+1PhtwFUBvfQ5x/RizvdQy/VZzPWaz3b8uST+wFWDABqKAC6m03uGx2X+xGKBxmL2iWDQjb777/Nt0Pf+xmlCyPT+v032hVv+yXD3lGx2wlr99/hJyvX1gFV20vf1e1DCaqr6xdz+wlzJe7T2jmWFACT3i7v+0IP9u0R8xF7Oh7r+zHn/+WLHc7CQGUP3vtePEUT0i7qHADORIUyEAC2BAACf0Hr3j73zhLX7x96RGkrMgbePFEiRDUeOA0P4wNiVAEaPBkTzh7f9yeCPCkSLADmRAESYAEf+/f6CAET4xdz4wtnzibiKAEH/9GH///+OAEP4w9pHeYEoAAAAFXRSTlP9+PLorFm9i97MnGlJOnoNKxwBAwB644ahAACClklEQVR42uydf0xUZ77/rQLyS4TBk4Yh2XTdxLrurr3qVcyuRdzvt7CSRYSrLMS7S/R6WZqy9Mem2dt+myLG2E4VOnXipLWYGnXStTXx/lOUq/QSgwgxWpEgFccTkE0cYZyJZ2YcyeDj9/15nvODmQG1OpR1t+8znBlmEPW85vPz+ZwzM+b9A8o0b85c7J5K/SMCMc2bG/cDkL8jZc5LmTnvByB/N8qcl/lM4rynVf+QQGYHEp5WA/nHA2LKnBfv9yQDy9OpfzQg4JHKmCftByB/H8o0gYc/5In/AchjKTMz1uYxL83vlxhL/wHIY/HALbY4MhICjDH/M09tGQIg04gjIw0HzhQ7HJlpIY+kuCVPwlNrIAAynUQSslNxFwvXB6xzU2YF/JJTYu7A0xvTpxnInIAnLoMi8RPIZOKRKD1tJnDITrPVEfLP+QHI4/WcMmf5A+ZUHNMngQFlpCRIwOF2SOb1FZJ/5twfgDx2TR3ye2bj0eOwEDAy01MSQoEAAw7ZbKnoX292xD29PKYTCL27Uz1MYoE42Mp3+4PEQsBInsUCMA5Jdihmtn5n/1hl0PX0ViEAMr0571yzn7FQIAk+5pFRmMTDjNS0OA6DSYrT4ZTMLL+xf2xsp9s/4yn2WNNpIXMyKM/ySCDiSXz4MRQuijQ3PT4tLjHkGRkZdkIOJwuZg9b1jf3gQR7r6W1kTScQpFjPxKOwDgAIUzyzQOQRYMydkzI7YZbidDkcTkUKQWazOeitXF9hBw5SZciZOi9T6GnkMp0uKymQkDHHAxzMLwcSM+eZJqWhskhOMivDDscIC5ohyW+x1lTm56+v2DkGGlz9FSEnynRdgPKUlezTCSTe4zGnZfsVtmWT34kljAmJgAbBSEuYyQIjLngnkLBUgkJFY6PdbgcKaExTfz48Fnxaanx8fGp6hlY1fvcU7p8SSEY288NAnP6XtlpgI7MncFq8Ao+Pm+kPBEacTsBglfkVjWQRYSB0Ho2S7ElOfoZ5SMycmJySnvndmWQCyT8fECghEGKMyf4tW/fLTKGeeXR/ak7yMx5kUk7UfCFrPrwTBzGZYCDMjx/nAm7UJ7OSU+cK5/XoPOaiWP3nA5IpIjqAbNq69SXcmTPCDwPhSPBTBU7GYcmvsE/KQosgQSdjjKcJsux0yrLiBh/PTJ2J6dHsIwHW+s8IZI6faUC2boHTQoUd5q0y4uDSQgwVeKiyYueDYWgpFv1CpGBMQg4mSRJz4puRgH9mcnyGFuVND6lWM+MCqdMHxDRtiQjGp2YSEYVZtkIWOC3jONCDlFCA4Yg6zVJ+xdhDcYgahCmoEXkGVlNZWWN1MkqLJRlQPNkJaZif06mYJsvn5swKZGdMs8uanozdhCAiERDHfgB5yS97jNyX3qcev4TjGwrmNz6chojoXtnBzKGa/IqKnTvtY2P2nTsbKyrWV1qAKCQ7XDAUQMkwahsio24kXq7G+f2epOkM6nNS00X8nJbeIoB4CQi0CU4rBU+qKVgivYZ3eGUFcDwKD7tVAg6ZQg2JPyce7axYn29FNe/m3suclJwyJwO2EiVk14nkJOmfMW1AZgf8ocS0dEIyHZUIB8L2b91PcV2hxrmJv5L+jAcZmCPE1o89Eg6o0mzBj0eHGsEFUCoZvJfT6UXu5Q/NSkpOS0mdk56eQUqfk4oWALJranb6zXOnEUjIT9mhFJf6/TsuHHYGuQkIEdnEhImQfZjBQ3GaKxsfFQdKQgt+fFJrIktpXF/p50wU/J+BJeDxM958CYVwDHifUmJM4hXRtAHhh8QpBzxJQPK9hnfR7NUsZP9+nvrOQhRBtJ8FHrJsflTzgNaDx8N+nDPJl4kJoLhDEtUqQgwsFIDiSp9OINl+5pWC+CcGPAlzvm8kpkSPZADZryZaJl4xInOteGTzGMsnHv0TEIgQjyj5Fkq9FIfDAS4K/n63V0Z6zJCOMTatHgtAkjyKZC2qZsGQ0+OPywCS7zvNEkBwQy3i5LVICp51St+Bh70SPPLDn7QXFzeiyRItNaBYmRkiK5GBQUK/0lJdJDPMrCRN5+j8jLQAgDQ1lVR7g4ocCM2e+71Fd2KfHJB42svFwzpKgAzENVnYR7/dbhzIyfNdazCKB0Rp785JA4q9sWL9+spKi5ehfvRb0cFvrK0J4u+d3pmVGRlw1sG8pt7akupg0KkEZsZ/f9Edea8BhCQy37iApCiCR/3KMV1IZyd2V+slyYn4Ea0HV5MiNwazRlgSUmXBAy4rfVqB4L+vSI6S8q4C2/wahBK/JyH9+/JbwjmhUhc4hM9KTkeEdZo5j6pj4wwk34qKJBpHRaUZgz/48cdTv6qdl2urqRM2zR6LCkMGE6mxld4r7bXNPxCUFQ+D38o0fT+FCIDIbJNmITzPSvJIDgRo4rG1qn9c59AczN/ZH+mt8s0wD3r+iQQcnAcUmN4RiRncjzuD1SByr7QDfktyugOzUr8PIzFpQLbs1+RAiEevPVQ5BtVvXRt27GE3DlhCmLdiqAa9FXj0BLI3AofKQ/IkmmAg0wqEYqgTYQRE7pRfnm9FcPfQQOHUB3eMAfEFKhXIMQQR4IEPk/k7fu3++sjWiES1hoGo0ux0mCvxw/A8IohDD44d/WHwKIi0A8fl2hrw0AxE6LGWDU1PDEQsS3iV4Hxb+b1797oKSuZLZCShFCpKvhcg7CURQVaupSDCmEMNIFtXRjbXZVkcf+H8K1jQInorKC4acVw1XR6X79p7oStXOjpOFyMPJrW341UVRiOncfnyaqvgIXmSI0aGv5ue3Lpm0G9AeaZIbKOtoAtISq/ASJjbH0iKSXA3PRSIIoAc21q1EomvU3NY9tZj9ZHNEU7L0tjPDQbfar0VHFeoVkfSXFvbvnNcRcKhgMrp06cuXqQVYA6jkf4I12IMaTMuf3JaMjQ7LX5Ohuk7M8ESzpOupMzQFopkybnRdrm87c69e+XNRU4qkGJjJCbTg4EozLlfBVJFlYhby7D27x6LBgIiEsPrqD3MliCZB+GALpPw5odxUOs9SnZV3LM1EkFDRUF4SlWegJAHPVexqGV6ZCgmeJuUJwfCnVYIRCy1NlvTabyXmppqeAcuBkaSiXaM6YFZFpIsDUj9VkR1p1ri7T5GKVa9PRIIbCi0Hu5KcqjmcfE4VyNh6NfE/Zi4x00vE8EL7MIFdyUrDHK7FS6v4nZLErHxhBJStIXGR5vb8MfAQuiAJVHzKGjdWNtks9lK8qySU0Jjx+2R0shIniRKxKHufWAd4kSSdQwbgIwdQ1TnEQQkWhtoiWlllIUoblTTlRILSdw87HBCqheiw32xcUIdP86tSLMjfOlWVY2Q6aUleKIiCTE8gRV5rxtN4FBSSvojjXjBPuYmBubEAAhvd/sl2Ii7prq6usYVlKSgdXFNUIGRJFLH8UkGGSbP60U+ISOEgAeA1I8hqushvTOHEi1YSTgQ77PZI4qMtXO1M3/61KliO0X1i7CUSLVjK2sva48QojrEjWWxIyg7QTgISQztRovV4nCIJyTJjTEJLJ9IibNT52qrjJPi4ItqofQnDOszxs8byCwIcRx5JbaSIgmG4mFPYiS0BkVWPHnrBOuFZB6EpL7fADKWcxIhvaovIoaAx8HDA8MIJBaRbBWfLraTnTRe1HScfxlQeFDhAI7rQI5z8wAOpC9OHHtmrSkqWrx49Wrhw6DFi4uqayyMDgagYBDpmYQUbUU+OiUWnOLNWG3MiIGF6E0M5oXlopNRndfc3lt+xbYRySAZyeNHEoRtP/ULMydzaAGUhceOcR7H6vtXbt3MHKKN/hWF9FYYSDiQ4RNvHjo84II153MQHcU8PFzUJXCoTODL9LiCn+ZMDAGHhKZyTdFigAChKIFLdY0TUGAp3gBfkU/N0BkY4t+nJiEP4EueMQAiGq8hxtD53bixuQllO1WJNl64hx7fSExkeqEHdIcSPahCAISYdNoB5CUViP1oFQzkZNSQz3DPO7cPfdTiYsh4QaT4Cn8FceQUtjABB5+qE1lVI32PO/gwAWV1kQU+KsgqF7dr4YXvoqksLrLCUshQ/AGE+cS4NAypAoKhzIzU2bM8yOAZH4+IERAK7JT9Vjc1NbV3tN1D/tvWa8s7EHR6Kd16nMLdJHxhaJKGNjUJ0OolGPTVNx5I/Vd2JFo5/eGVukxA7t749KPrLuYgIlfs3G+dIgHJzmIDSCO9BEcFFFEqW1zpJd9sLVrdLnjwndgmglLEDQXug09D0pBqUtzstLSUlJS0tOQEnBpB43xUV2I5JxZAoEyxbopGY3Wz7VT5ZZTtVJJQV57JIU/osYxkrhqcJHS0Jw4hTngsof1rEcl1IFU5gHLUHtY6sQggN27d/hRxhBPppRd2nlI1tvOUbinFO4vVuKI7MIOGg4bnrdWLEUzCdVm9TWQo1RYwUZCMSRKw8DV5be9H/Hcib8ZaSsyAiMkCxBE0GmttBaUdvV2EpA1deSWoSI9Tk2CIIY4DmfjEcbzMy3QVCBol4yykqp5DCW8uMoUDuX/r9qHhZwWRfsNAisGDK8p7CSoQYBTVSHx+vmhDmZqGRSOZDMriaqukQqFCxdiwAsyCCpB44mMHhA5QyK8TuVN+uhxI0Eo5XlIjKeyxCneTOS6VWwgtlUfxyEz0UEg3gIwZQKggzFENxDj1ww0ge765DyKfZJG34zWk/TQ6Iqe5fYxXZFQBjOexbCs5avIXA4ahaCrERduinZeDoDCFZoe5cIcYw6qLPJI7lBGjGKIlv0wncqqtq7y3vPQO7zf6mZtJ/u8aSUwmxKVZjORJnHj9FkWIDqRqPBBiEdHKqqAyZPj6PgABkfdaRtycSH8xgACJfYwwRDDRoZQtzn/++fzFi8vKhKlMJI1LNJzIKF9tlUEFUvfOmqLVq1cEFf7fjB0QE09SmUqkpABICnoLCgpO5QUVJMSK9F1bwCaktX5GksiWo6t03UAgWoxaqwKZdHKXgNwnjd7+S88wIyIqkGJxX1xs8CiGxnswPBzvvjQqZWUbNmyo49qwQWeiP7sB1QnHEkEFWKproOoinjavPjIisRhMdM2IrOP8TCXSbLN1lKMYsZXMd0hOairI0ndpARu5tNsbHUWIPV6AgWhCXYgVEAJSMTGQfGEhH98XuvFOj4sRPgFkbCfngb1wYMU7kWMRj0njynHuxzbUHSH96Ec/OoKDX0ZmAhhnL3195sy58+ePcOEVqlSwTSLgOFITxNranBgDMXEi3GvV5K2upb4WnGYQ37L5eY6g020U7o8+VeLms3AM7i4iXjEjxYJQF9o7j00GBKqUCEiLCmT0m2sHkWqhN/xcB5EYO63viYi9385pPDCslAGGKmKh62KZrg1cZCiiBUZbNA78nhWSmztmU0yBCBuhOCJL1E/AQD8LSohi1SUdtpLKoFeW1Jrk0adK4OktzGvUIkbzTGaWcTxQF9bvBxCnFD0+ypNbi6S42cjAwVt3CQilWlkgjb4WiMBVGW6LwNiLBZgHhJUN9P4/r9GguBId6R8eUlYTDeLB3CLHijEQk07EyXg2IaOCss7vxQzERRs1t1jokdMt4bLQXsfsghjmMHjMRIbNI3onNsLByxAAwbM7JwTSKCler3vEt/fLu6Mqkdd3ueitM5Lb0VHcAdnHNB4CDIhMRIWCCXBwHTlft+G48GBim0hEQTWWOtKRuvP4Eg+FaiS3m0YeYw5EeHdGRJhbcRISR/X806X3IBoUoqUDGEnCI6y4G4nUlk5aKjf6oOCBGlQ4LEGDA0GSBSB+yaplu1FJljubjTx7+MtvAITrm30URpyhGg4EBkJ7bAKHweRUFJWyc+fPnQOPcxvKLhqKhgKnRRAA7kFa4Qh63XyACIo5EOo/wQoYNbaq8/I2NjfbLtP8AxLg0nbe3HrUmkSMijrZ5s7NcESiDW/wIIfVaWz7V9JUw7HNfnPlpKdHtbx10BdyfaoBGb3xqStLARHz8x3CQPgdoNAWrjAqFwvPQefPFZZxY4kWmABEoQ6i7gE0FgIH6rRYTc3PmLjCNhMRRappbmpqtp0uKCmnZiM3Eortj1iTiNV6OKZOOCLxzzV4KIyt6+zUeBCQ/noYyxZ/MH+SJCvoOPHKWyeG/Z/eABA19339hIsRkeeugAN6v7pO020iJNjlnoHOnQENUllkqIdRFJ7TTOJ8YVlZO+/IA4p4EnvDNlZY/UEJ5kHypE4REJN6zJiMNcQmWEhpqY17LTKS2hJkXW6m1iQPP4kQYXdd5zoMMdPFjU06Dzc4GTzwCECq9nce2+QPvUhAqqLWxGukbG/bG0MuZgBBqrWHqhE5ZO0AkTEeSMKgRFMhHl9/faYwV/VdG8rGZ2Aai/OqCnNz8RR8FsSfFUYBrVhRY3UoEtXsXsbFW++xBmJ0GpMCPPu1gIjNVtrVa0NcJ+FREXPzwj2BLgb3QCDpfC4Rh3sTUzx06oeI52BNjox0TLvtr+pfCyAWv7SOL91GJVl+ZLz33h464DtkAIHTGmFe1USKyWNFIYmCkvv1pa8v5ZLBnDpdVld30Yj1ZYARIc1ScAcYCwHBinMYUFepS4xeRViH7rGmCAhsBEU2OQPJCSJNtnJqxV8sKC8tLT9la7aQz3y4kdBQiTjyFNUpBzFp3hABXeDgewEEVQiZkoMDWWuPOgX9wNAb9/44dKDlk9s6EGRafyGnpZCJFI/RpA+2SCphtpJ76RLHQUQ2HCnEXqNx/sFCrJAIAsQYSCiU9LHxSp86IPw4z+YnJUsyiNTa2su72sqvtDfbSmxNG8HD7UTe5H/gFRPFNIvM1qlA0HiL4GEwAZB6eCwEG7aJxrHs2EXGdMvQH++9NnSg5z0A0XX31sGBEWEixXbw0LZoWxG3XMIhdLHuRxvAgzayDZ52IQlGq6QWQtAYj+PIChkLh4AAChCLUigQBwcwdUCISIqHTz74OZGSZthHeUFtbVOzVVJkhmyLMj0znb0weW/Ro8Bj9XEgtCiC/I13L/GkJt1x1a8UQDZzIPBgkY2TA7teuffK0IETYUDQ973Oo0jNleLiK0Lgoe6joZw9m4s91FF2BPbBlVsIGoSjcMNlodrVPNc1tAK9CoVLs4xopU8tEFGQUHqKmcYmIGmykZqaixySjH8eOioyDjKMZLJReVqcIo/V19fXaWEwqvR56dkGj0gkVX3AsoWxlwhIPVqNkY0TOeuDe++2hAOBbuwBEW4iGhCNRhQT8MgVYMbz2HCOC/kUL82R70aFEpynAQ5SUAg+C7YSQYWfrjq1QOiImgWRaipHYMcb86odfE6jhoZSYCQUSQwjiR5hgHNbByDrnGgkhuZyHm7yV9HqW3uMuzZLpwCS0x/ROGEDB9vufYAph7+EARm9fcjnddMSu04jnMt4LLmLNDTgcV5ElUJYB6lM9BrPEA0OiDYhmj+nFtIKUg2yqxGJsCCiM0388ximFggkJo2Y4sb6Cy4WZrUwPHCwoDOvpL28gK+3uymSJMNIMqP/sMix+ggIw8zZMykWjxu9J+LRh20CrXPgxWNVAII6MaJx4ujZgwXMA+zEnwAkjMibqNfdSmhRbySOKFvReeQib8rltlJ47sw53M6cQlABDUEHOAQRvncEgYNZFxIpNcBTviVLFFWEnfiTqFCeciA0i5dA/VpqbEFMdsA6lOpmhHhab+ftRjLXmakEYIK+CWZKCAhahl5m2YSiRFbIh4HIhFBegscSQI6tjGycUJJ1r+2jUM+bEUBuHELSQ+V67xXaJpTO5IrgQSWGeHSOasRzhXiE8j1aQIBZKInVLMRjVWhhqVQcYCLMJJREA4VTDESE9uQAg9wK3hBiUGOj7Thw8KEU0W7UjMQUZSAKs/SR0DpBMYkb8IFQp3qLYrIZ0UUA6YwAsh5A3kYVdDArEgiIvHnCRcUhaEzCpCMcyml6r5/WeIBI4RVYCgyFbxHCAGfQauCoO798wYLl43JhFUnAzzPOKQYi0l8Pg7yY0qguQiyhZgofgKChFGo3Up0UZSQiglBIJyHJIqbAIoOHhsPYNC6b2GYVSF/E6Hs+gLx2707XvgFMOUQCOZTlBmvzc70qEXE/OZQzOLy53FbOcH19JRc0IGISAWUhHMMKbioQqKxac+HCmh07FhgJmFUCEjoE2RRMpxoIua0U9ONpPKiJ1Nzc3HTxcrPopIiT4JiXjMQ/O5O/Q4wUSzOQVoQGRcytWxDiDRlcxG4d9bf6OJDW3ZGNk8ET7+Jv3NMyQGu4EUSweMh9FteV8dtEKqRjzk3lLDBAubgXwgsRSBxmx0LdfdUtX3N1zarlsC986VpoDdJSjXqN7ikFYvTjQ3xgS1NBga23jSO5U16bx8gdSeFz2ZmoQchAWvugF8lAeLpr8CAE4kZfAspLCOkCiL1hZXiS5WSK7wP8fW/1PPvx6N0IIKhFZOGzoIdS4WEjlywGXS1hIRoPHQsHQ2hWmK06HfDo7l51rk4ULbjpguNiCmWcCaaYjQE9PP11gggMBEIrpZSaWzARclvVyEKgkMdP7xCTfqEfMpDW1ta+VoSQSB7RdoL9FuTIAFJPQHLCY7rkePZwmwDClwwjkHx8XeZ5VpgmYXLpDMRLx6+5ztAWqXPq5tB5AMLym8vO1dGDOohnxwIL9iuYpFC5TtnWVAMR3d8AEakBEa6S3rbSAltzR2/vqdqmoqCTL5tL+ly2uFgcN5BWiIcQGf6qNRzFOgDqNLisY1twL4AcjQCCxZCPuwjIieyD1yKBiDY8asMwGmFUDCy82XtWPFKJiNtEWmgdH1WWLasThrIKWg4uBpNzCx1BcluBBEFkCoEY3V8QsapELpbYyru6Sgt6OxBWyEKwwW9pMxBYTOdrt0QDIYRwyUwhHvh2HBDHS4aNUHK8jrsuamP9rSoyyUIZAiBvnFA+0pcMjS78p4PZXgoik8gwlN5FHEEueOCRISDZACaRWBYuNMylbil4kFZdENoBKPBfEDcSjJ0QkTi8H6cYiNH9BZGNRIRcVomtt7y8vLe2aaOXeZ0Sho1hs3wGQjtrkb3YysWTXvFdn2CicdmEJNjQpi3inoD8T31kJwtlCIi82SMP05JhdFgfRstXJzA5FXTfoSvQWfThDW24XKh7sIkEG1kqsrAdFwyteXn5ecEEOwyeEJFkIjKlQIz0V3IGHUSkqbb8Dpq/aP2i92uVnLgUxEZRkjBx8kIiRfQtIKB5LHJfXAaNVgLCNo+zl3XcXFqp855DQAxVovn+ewDperNn2I8FkSif9V7PsNctIYg8RIsIx9e5ILMIZAwkZZzHGd2DReHQU7CXdftYsGDBf+24sGbVGbXXIohInMiUA6G/Io2ISHIeShGbrR0jje3NtRurJRoTmm9rsi2mwS1uJAnJiOjkorjWKcyrui8KKfo9txCU8poP26L6r7VEoApUDNkrQ1QXonUCS8iKBgKfdTjbi0qk4GFAci+dBZFFAHKWAxG7wuPtun1ADwj2ggcoLMRsHWnhDrCBmRCSM5gmVYlkTjEQY0AIbR0r6kP040tKNmKUPOjEE5hLQdnezBvAFEn45QBePNnaSrcX8VhhRMewEbEjIF6mhpZ1wlZaO3kFUh85k3Vg6I8A8jkqjoFoIDSAcn3YG7JWCj0PPffcc4sWLYo0kLMkAOGmohEpvHjxUlg8mZSJGj9W8WXeBRCwLODmco4jWWHm2W8ciEw1EDH6i3EUDMvxXqPVOiwF6WxdR14HT4DLm2y8AQzxC5ichAjIFiZzhxWtPgv1Gh0gAgl/tXs3dbGi128B5DUA+eLNnsHrtGQYoVGsHA5jHC8rm2RWlc0c1kqQMbjknuXCIxFLBJTTp7ATdAwoYVzCeKxZXodEa8ca/njHAk6EPFedIAIb4bnWVAMxLhaqOGW3WMkEGmtRra1cNFLacFaiRdQkALIZNAjJOrIP5zqgiZRIvyjW6OG+s2r3/pwJh+QGh17BX/LFn04MiiXD6NpwmAGvkFN2kmRviNMJOWqAhagIHrn06JKujtOChrFFYsFuOaVgy/mhP1O3/GXQ0LTjiDAblPB18FqCCH1m0NQDMan9eEimJX6ntRoBBVVJgd7bQgNYEc1IHOeTpBe5gQCOuo2TqODh3ICDbnBX9a24XsBE13QHEPQWYQgHdr0eDYRavm5qX3oV2eUaNHQA2+AwIyzM+vxzyzkQ7rGIiIjwi8SDSakQkCVLyH0RhjWXLvE4YmjNAg0PIVkR9HqR2sxEsmmaYiDiLBv048UM3UY6PZRajRd7S47rvS00gBXVRiwvEpAtzEkGAomYYkBRK3iFnBbUByBV9Z3IekmRdaELrSwBhC+IREX1Lz/KQvLgkrMHek4MjdeJHgzXDR444HKHYCuumueXa0AEkrOL+P1DoCwlKgTi6pKl2EcQMTLhl5fXraC5Ofh2ajVOORDRj9fODVXLdltpeUkJpoS4yldTo83N3/mIIw3rkPJq8QQ32nRTQXTxgp6REp8cyzl20o6IHg1EbgGQttuv9wzSSW3RwrCD4s3eNTB4cM9bb7z99h+F3n77jbf2HBxuAZcBGcaiEBTH84gqFNs5FNxHKxLKWTKQVReuXr2wdCmgiC0Ky4JVOwjJQh+ISKLVaJpaIEZBImttFAjrI+j6llwpL+3qaitdbHZibVMhG6FuCXIs1ISEIxwKYbEwOQsG4hUZMhnI2FpKslbaJwPyzSGfa4D3TqLbWQOugY9febcN3jNCXW3vvvb2Wx9nDw1d9w4OuryAwmAolwQTgvJwJstxt+YqKKwBFUEjmglKkyMLdlx4eYHDhzelelrm1AIxxlGcBpHaEgz/lhY0l5Q0X7bZcCWwEYmJbAs7CzXhQULgCLsHK9//2eLTq0YYSH1fJ0ZK10bN9WpA7sMzMT5MGqV9A4M9ewSBO+OlY/ngtd/vY0NDPkQWJ4KKZIXz4kj0bWIsuBWuWvV14ctXhS7oW7StILJTJrZj20A2uS0pzYTAa5pKIMY5JFobBbO/iOW9sI7S8g4bOo0Hgrvmr+brVoxWbN1qBmxAETfKhwc//O2/gpjCy8aTnTn9K5Fn4TInkwK5dvudHkqzbk1gIQDyFrGIshCDy50PXnvj4K6hAdegS3GHKMyrZjI5lVX4Kly15lLhEp0F7g0uXGGmsorWsLa7WpjC0G2ND7/oFh7HGggk5uNFGwVqKm8rP267jI52Owr34V0fLjplq6XT25m6TLi54eQEanUw+d9/8VOst4PZOnJj9npkvrjiyaQW0nUN6S0MIdpCvkEMGeTtrsmksep65fcHuZ3IxMRZs+Cs0MRUXl4CC1lyYVXh18uu6rqg3chGwiSsZA0RepW1hBS3x5MUTwgyVZnwMPZAqCDBhC6fMyXZCsrLT9HJiNX+oPnDN0ppSt62mNwWScE6CIg0YNOke6zffPbTdYMobKhqwartyr7WtRgBWtkfDcTVQ0BGyRJkRrV6dJY1yLsr0EOggMkbHw3tcg8OKyOSOaSayURMCl9+uRB3a7iBhMlA8s7291/dtm3bq+9v3w4YOhVCMtLiY16PBx8+YVhIehrOrYk5EHFpKEmmK7fQyBa1UfKK6Kx63/Pvivz3sm2j6rbIUjZzHuFUNjHXv/3us5+uRa7FvVpnlb3vZN/K/pUYkosO6qIOuYtJRcqzbo9G1SHZyuAu1I58e7ihtL32lm9oABFeUWAmK3g0CbcUerjkwhJQWXp1aeElMhBDO9T7d7a5fAOQ2GUPH972/naNCpC8fzirZSDLH2Cz4tLiU1PxkRiJEg3UxRiIUZAo3qBoo9DIFuSo3thULv7PMJImuC1BxM2NpMGAgj2qwuF//wWAoBqBEcFC6nM6sbPTpBzJHgHkNfxeWAZmTAavI4pEtXtdjH0QefCxTcrkg7cPDvW4XF70gsyS5rnCoVxdxvdXzwLKeK1aJYxkm28g242cEqIdCz0LLsy191XDWN7f68rOYnSBFFIgNGcqXJZRkLj1q39hsK8mj2ZOS/SRlFpbnoyKBPI6GQIJOOhUyECGfQdeABCQoUSs9Wg94krDWM5+dfK9ym4AkWThj764dR9n36KPeAhEwoD8CU2ufW1dbVxd0J0HchGu67U9PSfkQa9bdsNzaUgMLEuvLrlELP7fJTKQ7qvdawSPNcsFj8MtXkpvDbkJDMuCvWS79pKxcG3fxuhCHCHJ48fcbeyDukEE/XhRaWM+CIU7rxSbLtps5eQ07vC5LTSDuXAHIyEaKhREkGF4rHFAdq9sPdmXM3ZS7S72r6zv15uLijQ49BYO4RfX7uKckL3X3a5Pbt8fjUiyMGR6bfQWNDo6evfujdtffPH5520gY2CJZvLuG74h96DCRuC5gCTcTLqX0X5Z9xKKIERkFXZ4sFwYyN6WETaROJVs30DLQLZ/+DA0zCQqFP1J2qiUiRRrIPitvCCBUwIKvW7H2lVvyWVMNpIK5gOIomhG8iKRgAgKst0ReCwDCJ4jt1VPia8AMi62V4Ycvo8w5/I5asJbNz79+EQ2e/3+7dFb41pZw1mHxy0ljo4CzTVo9O5t4qJiCWciPNfv/UMhIFHkcCQwkO6llxBHgOXSsu5uGMjLC7qJyLLlHM/2KB7RxsKys3y+Z7P9noAnO45wkFQUJvFNzICIFRLmlxTJKNohW1NpKZKsK8i8Ck5ZJcXJVxL1Vgpg0BcFjpbffPbnz/6VgGD9SmTClGidVDsnRn3YjwUqp4uietstfvivvTnQ07Pv0I3bd0dH1SmH68quqC7wXWxAAzC3gAVUJmQCJN6hbBVJzXKDCExDYClc0s21fFU3oVi69Go3to98XvYw4VQSL2CIK87qEmMHqsAkRkBEQZLtCcna7ANUixQYzd87pQUduBAEDARnxSHb0ozEASMh0fT18Ie/1YCg29WwGzbSVzXWAChCdpDReMDQUGVQmjVKx//ujU/2ZP9haN97X96+ISxkzx9cryPiTyLicgtUdCigYiARVpIlkIQMJDeXcSw3lxQu5TxePgJDgaUsXUYGciHbzR5B/lmzU9PnqpcPy5gTnzY7LiEpcVZiYmJSQnJaPF4iJjEBYhQk4lTEZlDh51CXdrTTyW8F7XlMBqwSUSQaRgLRUmHW8AsA8pOTMBZ8PnFOTisMx17fYExk7YbvEjzosjOKmyblENW5Q7rxzaG/7Btu2fvmoftch977FGgeKFABlPu3CUo0kreHYSVe+Fcze17wWPLtUrr79leAIgzkHL9bRnjgsQYU9kiaGYdP6ctIxYePzwz51cueIaBgj4f0agY5mxgBMYkBIX7im2YiTTjTqpcuk1KNFuSz8y+in5KnuS2UgBYYCRiw4Q9/Bx4viG8srfUrWxtac/qrGhqqNEdVtdv4xGEIJiKCCAnp7+37X3566NCn/Ft8h/jxSIIDg6WImIKdjqQNVqIM0jy42bFAAFnC95rHWla3qlvX1Zvv60AkbPpdtPz0KX0h7D1+9ackCjFathwIZMfRyGcsgBgFCUpELFVBzdgwilJbm1fDgjixZ3V5Fy6maRO9LRIlwC8yeCy3jwzkP082bGEOtnn32O6TDa31/TkNR/Xmu/1v9QYPZYS5UPh1gYcQosc3N27c+Eb77hF5wFLIfY3CUO6EI4Hj8u0adjG3GkqW/kpg0TzWqjoYii4NiMQoYnv82OOIT8hEAhS8IqnBXuEFi29AVZbbi8n5dPit2AAhY0v2MJp9oDWrZmg1P9EKOKx56kXp9N6W6DeCBxv+JQzkz0iyyH2xdTn9cFm4lnXO0a+MgrA+R7tMFvr0AyOy70AbVSKQTgEcHkd3yXvd0JgYSN7o6RkcBnzyW0teJyBL/xdfOP7d3WeWd09gIR7PMwm8Co9PS06ii2Ey6UGJlzsbJHxs+DDKFAhdl8M+n1t8GIUpBkD0FRKJzz5IXovFalHEmSQ1ec22AlEW3xG9LWHjbhn5oFtpgYFQTG+wMHgsDmS3fazqq69gFjqReuKB4Dk8sPfj68Ounn13uuCzYiIYisFER/LuW7taBmXyW9ZVSzkQYFkGHvBY4GLoJmIIjc8mxSMsa8qIT57pB5MJWCichVEwGj2x7XsHQm6mTXPNiN1n2ip0/X4JMCBmrd7IV9tPt6mXSTmO3pZKBNW7MJDPPvvsJw0/5xMqOf31JwlI/d/+lhPWxeKubiSbHbq297oL1eGdL4hIDJnc/rwrDMkr+4Z8B2DBWdLzHAiw3IS6l54BF9pUIO/4mM/LP24k/IrK6SlJEXWiMuINEYvBvYRCB4EvNIap+N+O3ySmgmMCRHxGDxUT2ElKjX5eD9Kuy3rZTpdJCTqNdwwiCIWQhqMvUghpyOm3H22gy8Lm/I1PkRolugwezPfJbVytt8d1YOitz2/AacVM5Lvuf2GYCf1j//gRinemeAesyxFDlpxdcpO0BHfd2gYgV92D21pGpOSMiM83TombaRiGMqKwZwdafMO8vXUVrs4QuvRXfwxRF/IwEUmOGZA0GAitjdNSOyiQmqPL9tJqHQhSrP/733/+jELI0c3MgRJ+N1CcbAAK+//89a9GF6sy6KC3GDUTR1GgDykHhvZ8eRuBI7ZM4Lq6xiFpe9s35KKz5EQGLICcXcrvuvVt28jVV3u8I6G4eLrQdebcjPT42QnPUE7FWZBhZKN9wg5v4ywEDAPH1R//7D9e+MXvfnNg2HX4/e2Hs90SnVkdAyDaNX+2bMp2eiV1FVEvEzu6SpEC99I1ORYxUBNSqK34GS8Lj9JUyr8cbajvr284Sd7KnvPXv/6VRnu1gO4ePoHBH1Ggt/QcaDn8ye27MJKYapTMhGjo0f0EorvCzDU6kCXLburqptuFlgs33/cNeEcCbvOsWbNmZjOkWzyb8pJh+GAYrr3vvwMWBgwDx8/+Awfgdy2/bGGMPdsychh3/FoXM2L0mc/O7H85CiJykHEiWpV4saOW1nYLjqMs+VAyDOSXv4WBgMhPjjZYgOnnR0/Spd5bgYWQ1OdA9eqVeplLzP2IAn1fy/UB359uUWcxtoKZwHMZSN7dg+USRTaPLNCA/ErHIfbfvr/925vowF9/1j0yMiKuueHlLLI0wwALQaNb34R+Ru/HF37zoTKsuCEFw7cSjaDGBkhKAB1Gx8+/Orol2+lEQWLYiK29tLS3g64027wYQ/J0AgWJ2rzggRCy9ujPqVvcIOLHylaEdVV2O04sDDmJxzvf4PBrBfon73x0/Q8fvXcDdXnMkVy78bmB5M5rB4cw5yXBSF6/efPbm9xjhTO5Sju+CtUy4PNlZWX7fD7BAobRTWZEAhWdhkDy4//87Ne//vMLv/k3l3qWe2JaAgv4JbqqaSyAxPHPLfrq6FdHN8NGWLCo2RjZwumIl2lxtybIUCZaQQ4xGg4LBiI8Fo1iWxqOHqWOiX13627Rczccluv6x9f0Ihxgbt/48r0/fTzyzqH7sdfd0Wt3KZioIb7rbd8J+C2z9c1vv715k1OZSN92X3jn/W17Dx92HUZ18ep2clLELEIEQyD59f+n7vxjmzrTfK/+oBTaaQsdZLmeqkLAHwgQSIO0C6jSTJI/UoVKV4quItGUDZOb7TZd1FKSAEOCMg2IH2EUJUpufiBBfm1BMcwmDR3ZrFZNlX82qNpo5Ei30mDHdkh85tiOk+PA2Fg93O/zvufkTXychuTYDvuYTTzT2ZT4c57fz/u8N8q27rpxI79Ymk/mX3lz0/rEo8T0y6sAstlwUdxL07BYxwDk22+P2TzYmomz7Ny1s7T9z//7QxxTDf398m/+hfGIKMV7wINk2234dA0IPDoRoWsqxK532rEohn64lozOzo7ae3riTzOKhLuSM6g5ehK2r4jIj9oLkuRNfqR/zD5tzZqllN/u2kVIoB6lW3PgPw6Chy6PXsCdzBvnpt/ZBCAmiIi7CG3vAwiQoCrioRTx8r+yntVvfvPhP6HVvgUz8h/++UNmshRUecGDWazbt+HTORAyWuTRkY38MwdCJZOQ4m0XvUFRJRllVd5MItFcyReNEz5Jddb+5Tt87PhjwKJHXuItxIgEJHaV0VeoRumBshs3DjcIHnTMCUR++foLj9atGAiynyQgr9DGV9s+6AcR2Qc6kQdgAgEI/gX7H/5MZ0MpI1Q5D81i3b5MQ8C3SSgvRPD7xz8CyR9o4t0TggM5pfEwQsmMxBEGAwlY8D/YKhj1KV7pv7+DHnAoqZgIxyJi48U8cm5s/Su+wpfnoEaxF7HY4nVbG9jkyIYVAsHWvsU3ZGxGVsgGS3aDBpN9NInILg94oIqsnY6G0n8tF25hPHSLtRsWjDSEhNcV/3AVSG5f/Qe6T0eiOd6sSxKSLzonOmVrFGYLotNYWle4FaPvi7Dk9Of8FV/xGDLvWVQoL64V06ZpFGHWrxAIGSjMTSTFWLAsAKLL7ss2Ft9G/n6Zeu2/4TPyjIdsayhi/nzeYu1D5xARGgNCnp0uXPtnQvI+eMjebtqanHWJk+EK8LyEZgPOoC4Ps/U3UhJBJTUUIx5mwUr7YbDAQ7MNew5yHthcCZnmTN5aRaZOn/+LfHJCWCwNiCCChIROLHyoe3bIvzIexVAPxqNft1jv2x6EIkgMIVxJyFqhzHv7I1YWpuGStREgYakiU5JPnE5f0Nn7V05EMEmpKgKG/o143Niq6wdJvjOoOfOX169/+aU5xLxsR+DmlQPZiOFI7I8Qm2BfmFMZkIfzQJCQWDEeuuVDHvsKHoX5e6tOngQOjgQW61sAUT0AoguzVn/4w//9X+wCsBPgsUYSZ6kiC7goTeye8MlRtZqICGFIlvb1gk1pf/8uslqcx8n+ww0yD3ffeJNfh7ox8YiryEqB0BpFrS4pdsxAPAkAWYAERCKKWMXxL5xH7smqk/39QLH3pGaxKA3hMbOQP0JIQZgDWVPhroTXRr+ckKWE86sfdSRGLktTAY/Sv/xlVz89jKV7hYLMrdMvEKWYFwdKVgyEr8wQp0z5ThNeOQEOINGoPERCovAzJODxG26vcm/0nwSO/v6dO8lw7b8KDUGiTuX320kCxQl6jSsbsu/d4Up0s+Wy+EJ+ciT0ErLYq6RgAntV07z117uo/dO/a5euIOz4NGDwgSDsB331EW02XikQVNrZ1huUN0WlF0Csx6AhDxkR/nrfpnrYiQUwebBFYem5zuP/0FfUsSCUh2Dvxu4kIFAQj5NnIGuMJOzu0wrBnwX9PsXZ+xdGIyUWPf5ajCWnqqt/16+3nrzRfOPG1q03ILnFpCCLLtwkJBvpIt0VAYGBWkdBLjrzr9L6NN6b4hryu4caCV1HkJB42IkFJCDMYO2BvSLZ+yumIOUIk6EiCLNClFYmK4hUTyXFNRcEXLBb3Gxd6Z7oDEbnkJFoPJa0X4LIj4d2VbVW5RyC2cLzuPW/6aHcUygbNmHzSdA3Vgxk4zsAokRVz/QLo...
[truncated]

@llvmbot
Copy link
Member

llvmbot commented Nov 20, 2025

@llvm/pr-subscribers-mlir

Author: Tuomas Kärnä (tkarna)

Changes

set_op_layout_attr and set_desc_layout transform ops wrap xegpu.layout in an xegpu.slice attribute if slice_dims argument is set. This is needed for annotating vector.broadcast operations, for example.

Full diff: https://github.com/llvm/llvm-project/pull/168929.diff

5 Files Affected:

  • (modified) mlir/include/mlir/Dialect/XeGPU/TransformOps/XeGPUTransformOps.td (+17-7)
  • (modified) mlir/lib/Dialect/XeGPU/TransformOps/XeGPUTransformOps.cpp (+28-7)
  • (modified) mlir/python/mlir/dialects/transform/xegpu.py (+8)
  • (modified) mlir/test/Dialect/XeGPU/transform-ops.mlir (+38)
  • (modified) mlir/test/python/dialects/transform_xegpu_ext.py (+48-1)
diff --git a/mlir/include/mlir/Dialect/XeGPU/TransformOps/XeGPUTransformOps.td b/mlir/include/mlir/Dialect/XeGPU/TransformOps/XeGPUTransformOps.td
index 16044838aa27d..29579acc727ed 100644
--- a/mlir/include/mlir/Dialect/XeGPU/TransformOps/XeGPUTransformOps.td
+++ b/mlir/include/mlir/Dialect/XeGPU/TransformOps/XeGPUTransformOps.td
@@ -42,10 +42,12 @@ def SetDescLayoutOp : Op<Transform_Dialect, "xegpu.set_desc_layout", [
 
   let summary = "Set xegpu.layout attribute to a xegpu.create_nd_desc op result.";
   let description = [{
-    Given an `xegpu.create_nd_desc` operation, this transform adds `xegpu.layout`
-    attribute to the result tensor descriptor. The layout is defined by the
-    `sg_layout`, and `sg_data` and optional `inst_data` attributes. Returns a handle
-    to the transformed op.
+    Given an `xegpu.create_nd_desc` operation, this transform adds
+    `xegpu.layout` attribute to the result tensor descriptor. The layout is
+    defined by the `sg_layout`, and `sg_data` and optional `inst_data`
+    attributes. If `slice_dims` is provided, the `xegpu.layout` attribute is
+    wrapped in an `xegpu.slice<..., dims=slice_dims>` attribute. Returns a handle to
+    the transformed op.
   }];
 
   let arguments = (ins
@@ -55,7 +57,8 @@ def SetDescLayoutOp : Op<Transform_Dialect, "xegpu.set_desc_layout", [
                    Variadic<TransformAnyParamTypeOrAnyHandle>:$inst_data,
                    DefaultValuedOptionalAttr<DenseI64ArrayAttr, "{}">:$static_sg_layout,
                    DefaultValuedOptionalAttr<DenseI64ArrayAttr, "{}">:$static_sg_data,
-                   DefaultValuedOptionalAttr<DenseI64ArrayAttr, "{}">:$static_inst_data
+                   DefaultValuedOptionalAttr<DenseI64ArrayAttr, "{}">:$static_inst_data,
+                   DefaultValuedOptionalAttr<DenseI64ArrayAttr, "{}">:$slice_dims
                    );
 
   let results = (outs TransformHandleTypeInterface:$transformed);
@@ -63,7 +66,8 @@ def SetDescLayoutOp : Op<Transform_Dialect, "xegpu.set_desc_layout", [
     OpBuilder<(ins "Value":$target,
                    "ArrayRef<OpFoldResult>":$mixedSgLayout,
                    "ArrayRef<OpFoldResult>":$mixedSgData,
-                   "ArrayRef<OpFoldResult>":$mixedInstData
+                   "ArrayRef<OpFoldResult>":$mixedInstData,
+                   "ArrayRef<int64_t>":$sliceDims
                    )>,
   ];
 
@@ -72,6 +76,7 @@ def SetDescLayoutOp : Op<Transform_Dialect, "xegpu.set_desc_layout", [
     `sg_layout` `=` custom<DynamicIndexList>($sg_layout, $static_sg_layout)
     `sg_data` `=` custom<DynamicIndexList>($sg_data, $static_sg_data)
     (`inst_data` `=` custom<DynamicIndexList>($inst_data, $static_inst_data)^)?
+    (`slice_dims` `=` $slice_dims^)?
     attr-dict `:` functional-type(operands, results)
   }];
 
@@ -107,7 +112,9 @@ def SetOpLayoutAttrOp : Op<Transform_Dialect, "xegpu.set_op_layout_attr", [
     Sets the `xegpu.layout` attribute of an op. If `result=true`, sets the
     `layout_result_{index}`, otherwise `layout_operand_{index}` attribute. The
     target operand/result value is defined by the `index` argument. The layout
-    is defined by the `sg_layout`, `sg_data` and optional `inst_data` attributes.
+    is defined by the `sg_layout`, `sg_data` and optional `inst_data`
+    attributes. If `slice_dims` is provided, the `xegpu.layout` attribute is
+    wrapped in an `xegpu.slice<..., dims=slice_dims>` attribute.
   }];
 
   let arguments = (ins TransformHandleTypeInterface:$target,
@@ -118,6 +125,7 @@ def SetOpLayoutAttrOp : Op<Transform_Dialect, "xegpu.set_op_layout_attr", [
                    DefaultValuedOptionalAttr<DenseI64ArrayAttr, "{}">:$static_sg_layout,
                    DefaultValuedOptionalAttr<DenseI64ArrayAttr, "{}">:$static_sg_data,
                    DefaultValuedOptionalAttr<DenseI64ArrayAttr, "{}">:$static_inst_data,
+                   DefaultValuedOptionalAttr<DenseI64ArrayAttr, "{}">:$slice_dims,
                    DefaultValuedAttr<UnitAttr, "false">:$result
                    );
 
@@ -128,6 +136,7 @@ def SetOpLayoutAttrOp : Op<Transform_Dialect, "xegpu.set_op_layout_attr", [
                    "ArrayRef<OpFoldResult>":$mixedSgLayout,
                    "ArrayRef<OpFoldResult>":$mixedSgData,
                    "ArrayRef<OpFoldResult>":$mixedInstData,
+                   "ArrayRef<int64_t>":$sliceDims,
                    CArg<"bool", "false">:$result
                    )>,
   ];
@@ -137,6 +146,7 @@ def SetOpLayoutAttrOp : Op<Transform_Dialect, "xegpu.set_op_layout_attr", [
     `sg_layout` `=` custom<DynamicIndexList>($sg_layout, $static_sg_layout)
     `sg_data` `=` custom<DynamicIndexList>($sg_data, $static_sg_data)
     (`inst_data` `=` custom<DynamicIndexList>($inst_data, $static_inst_data)^)?
+    (`slice_dims` `=` $slice_dims^)?
     attr-dict `:` qualified(type(operands))
   }];
 
diff --git a/mlir/lib/Dialect/XeGPU/TransformOps/XeGPUTransformOps.cpp b/mlir/lib/Dialect/XeGPU/TransformOps/XeGPUTransformOps.cpp
index e301d4d9bd108..8995ab3082d24 100644
--- a/mlir/lib/Dialect/XeGPU/TransformOps/XeGPUTransformOps.cpp
+++ b/mlir/lib/Dialect/XeGPU/TransformOps/XeGPUTransformOps.cpp
@@ -167,7 +167,8 @@ getLayoutAttrFromOperands(MLIRContext *ctx, transform::TransformState &state,
 /// Replace xegpu.create_nd_desc op with a new one with the given layout.
 static xegpu::CreateNdDescOp
 setDescLayout(transform::TransformRewriter &rewriter,
-              xegpu::CreateNdDescOp descOp, xegpu::LayoutAttr layout) {
+              xegpu::CreateNdDescOp descOp,
+              xegpu::DistributeLayoutAttr layout) {
   assert(descOp.getMixedOffsets().size() == 0 &&
          "create desc op with offsets is not supported");
   auto oldTensorDesc = descOp.getType();
@@ -212,7 +213,8 @@ void transform::SetDescLayoutOp::build(OpBuilder &builder,
                                        OperationState &result, Value target,
                                        ArrayRef<OpFoldResult> mixedSgLayout,
                                        ArrayRef<OpFoldResult> mixedSgData,
-                                       ArrayRef<OpFoldResult> mixedInstData) {
+                                       ArrayRef<OpFoldResult> mixedInstData,
+                                       ArrayRef<int64_t> sliceDims) {
   SmallVector<int64_t> staticSgLayout, staticSgData, staticInstData;
   SmallVector<Value> dynamicSgLayout, dynamicSgData, dynamicInstData;
   dispatchIndexOpFoldResults(mixedSgLayout, dynamicSgLayout, staticSgLayout);
@@ -225,7 +227,8 @@ void transform::SetDescLayoutOp::build(OpBuilder &builder,
         /*inst_data=*/dynamicInstData,
         /*static_sg_layout=*/staticSgLayout,
         /*static_sg_data=*/staticSgData,
-        /*static_inst_data=*/staticInstData);
+        /*static_inst_data=*/staticInstData,
+        /*slice_dims=*/sliceDims);
 }
 
 DiagnosedSilenceableFailure
@@ -246,6 +249,14 @@ transform::SetDescLayoutOp::apply(transform::TransformRewriter &rewriter,
   if (!status.succeeded())
     return status;
 
+  xegpu::DistributeLayoutAttr layout = layoutAttr;
+  auto sliceDims = getSliceDims();
+  if (sliceDims.size() > 0) {
+    // Wrap layoutAttr in a slice attribute.
+    layout = xegpu::SliceAttr::get(
+        getContext(), layout, DenseI64ArrayAttr::get(getContext(), sliceDims));
+  }
+
   // For now only create_nd_desc op is supported.
   auto descOp = dyn_cast<xegpu::CreateNdDescOp>(target);
   if (!descOp) {
@@ -257,7 +268,7 @@ transform::SetDescLayoutOp::apply(transform::TransformRewriter &rewriter,
   }
 
   // Set layout attr in desc op's return type. Replaces old desc op.
-  auto newdescOp = setDescLayout(rewriter, descOp, layoutAttr);
+  auto newdescOp = setDescLayout(rewriter, descOp, layout);
 
   // Map result handles.
   results.set(cast<OpResult>(getTransformed()), {newdescOp.getOperation()});
@@ -278,7 +289,8 @@ void transform::SetDescLayoutOp::getEffects(
 void transform::SetOpLayoutAttrOp::build(
     OpBuilder &builder, OperationState &ostate, Value target, int64_t index,
     ArrayRef<OpFoldResult> mixedSgLayout, ArrayRef<OpFoldResult> mixedSgData,
-    ArrayRef<OpFoldResult> mixedInstData, bool result) {
+    ArrayRef<OpFoldResult> mixedInstData, ArrayRef<int64_t> sliceDims,
+    bool result) {
   SmallVector<int64_t> staticSgLayout, staticSgData, staticInstData;
   SmallVector<Value> dynamicSgLayout, dynamicSgData, dynamicInstData;
   dispatchIndexOpFoldResults(mixedSgLayout, dynamicSgLayout, staticSgLayout);
@@ -293,6 +305,7 @@ void transform::SetOpLayoutAttrOp::build(
         /*static_sg_layout=*/staticSgLayout,
         /*static_sg_data=*/staticSgData,
         /*static_inst_data=*/staticInstData,
+        /*slice_dims=*/sliceDims,
         /*result=*/result);
 }
 
@@ -326,11 +339,19 @@ transform::SetOpLayoutAttrOp::apply(transform::TransformRewriter &rewriter,
   if (!status.succeeded())
     return status;
 
+  xegpu::DistributeLayoutAttr layout = layoutAttr;
+  auto sliceDims = getSliceDims();
+  if (sliceDims.size() > 0) {
+    // Wrap layoutAttr in a slice attribute.
+    layout = xegpu::SliceAttr::get(
+        getContext(), layout, DenseI64ArrayAttr::get(getContext(), sliceDims));
+  }
+
   // Set layout attribute for the op result or operand
   if (resultTarget)
-    xegpu::setDistributeLayoutAttr(target->getResult(index), layoutAttr);
+    xegpu::setDistributeLayoutAttr(target->getResult(index), layout);
   else
-    xegpu::setDistributeLayoutAttr(target->getOpOperand(index), layoutAttr);
+    xegpu::setDistributeLayoutAttr(target->getOpOperand(index), layout);
   return DiagnosedSilenceableFailure::success();
 }
 
diff --git a/mlir/python/mlir/dialects/transform/xegpu.py b/mlir/python/mlir/dialects/transform/xegpu.py
index 7169b5e28ab5e..5aa6453b7cb8a 100644
--- a/mlir/python/mlir/dialects/transform/xegpu.py
+++ b/mlir/python/mlir/dialects/transform/xegpu.py
@@ -62,6 +62,7 @@ def __init__(
         sg_data: MixedValues,
         *,
         inst_data: Optional[MixedValues] = None,
+        slice_dims: Optional[MixedInt] = None,
         loc=None,
         ip=None,
     ):
@@ -92,6 +93,7 @@ def __init__(
             static_sg_layout=static_sg_layout,
             static_sg_data=static_sg_data,
             static_inst_data=static_inst_data,
+            slice_dims=slice_dims,
             loc=loc,
             ip=ip,
         )
@@ -103,6 +105,7 @@ def set_desc_layout(
     sg_data: MixedValues,
     *,
     inst_data: Optional[MixedValues] = None,
+    slice_dims: Optional[MixedInt] = None,
     loc=None,
     ip=None,
 ) -> OpResult:
@@ -111,6 +114,7 @@ def set_desc_layout(
         sg_layout,
         sg_data,
         inst_data=inst_data,
+        slice_dims=slice_dims,
         loc=loc,
         ip=ip,
     ).result
@@ -127,6 +131,7 @@ def __init__(
         sg_data: MixedValues,
         *,
         inst_data: Optional[MixedValues] = None,
+        slice_dims: Optional[MixedInt] = None,
         index: Optional[Union[int, Attribute]] = None,
         result: Optional[Union[bool, Attribute]] = None,
         loc=None,
@@ -156,6 +161,7 @@ def __init__(
             static_sg_layout=static_sg_layout,
             static_sg_data=static_sg_data,
             static_inst_data=static_inst_data,
+            slice_dims=slice_dims,
             index=index,
             result=result,
             loc=loc,
@@ -169,6 +175,7 @@ def set_op_layout_attr(
     sg_data: MixedValues,
     *,
     inst_data: Optional[MixedValues] = None,
+    slice_dims: Optional[MixedInt] = None,
     index: Optional[Union[int, Attribute]] = None,
     result: Optional[Union[bool, Attribute]] = None,
     loc=None,
@@ -179,6 +186,7 @@ def set_op_layout_attr(
         sg_layout,
         sg_data,
         inst_data=inst_data,
+        slice_dims=slice_dims,
         index=index,
         result=result,
         loc=loc,
diff --git a/mlir/test/Dialect/XeGPU/transform-ops.mlir b/mlir/test/Dialect/XeGPU/transform-ops.mlir
index ff0accdec7532..561034fb5880b 100644
--- a/mlir/test/Dialect/XeGPU/transform-ops.mlir
+++ b/mlir/test/Dialect/XeGPU/transform-ops.mlir
@@ -121,6 +121,25 @@ module attributes {transform.with_named_sequence} {
 
 // -----
 
+// CHECK-LABEL: @set_desc_layout_slice
+func.func @set_desc_layout_slice(%arg0: memref<4096xf16>) {
+  // CHECK: %[[V0:.+]] = xegpu.create_nd_tdesc %arg0
+  // CHECK-SAME: #xegpu.slice<#xegpu.layout<sg_layout = [8, 4], sg_data = [32, 32]>, dims = [0]>
+  %0 = xegpu.create_nd_tdesc %arg0 : memref<4096xf16> -> !xegpu.tensor_desc<256xf16>
+  return
+}
+
+module attributes {transform.with_named_sequence} {
+  transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
+    %0 = transform.structured.match ops{["xegpu.create_nd_tdesc"]} in %arg1 : (!transform.any_op) -> !transform.any_op
+    // CHECK: transform.xegpu.set_desc_layout %{{.*}}
+    %1 = transform.xegpu.set_desc_layout %0 sg_layout = [8, 4] sg_data = [32, 32] slice_dims = [0] : (!transform.any_op) -> !transform.any_op
+    transform.yield
+  }
+}
+
+// -----
+
 // CHECK-LABEL: @set_op_layout_attr_result_default_index
 func.func @set_op_layout_attr_result_default_index(%arg0: memref<4096x4096xf16>, %arg1: memref<4096x4096xf16>, %arg2: memref<4096x4096xf16>) {
   %0 = xegpu.create_nd_tdesc %arg0 : memref<4096x4096xf16> -> !xegpu.tensor_desc<256x32xf16>
@@ -212,6 +231,25 @@ module attributes {transform.with_named_sequence} {
 
 // -----
 
+// CHECK-LABEL: @set_op_layout_attr_result_slice
+func.func @set_op_layout_attr_result_slice(%arg0: vector<256xf16>) {
+  // CHECK: = arith.extf
+  // CHECK-SAME: {layout_result_0 = #xegpu.slice<#xegpu.layout<sg_layout = [8, 4], sg_data = [32, 64], inst_data = [8, 16]>, dims = [0]>}
+  %2 = arith.extf %arg0 : vector<256xf16> to vector<256xf32>
+  return
+}
+
+module attributes {transform.with_named_sequence} {
+  transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
+    %0 = transform.structured.match ops{["arith.extf"]} in %arg1 : (!transform.any_op) -> !transform.any_op
+    // CHECK: transform.xegpu.set_op_layout_attr %{{.*}}
+    transform.xegpu.set_op_layout_attr %0 result index = 0 sg_layout = [8, 4] sg_data = [32, 64] inst_data = [8, 16] slice_dims = [0] : !transform.any_op
+    transform.yield
+  }
+}
+
+// -----
+
 // CHECK-LABEL: @set_op_layout_attr_operand_minimal
 func.func @set_op_layout_attr_operand_minimal(%arg0: memref<4096x4096xf16>) {
   %0 = xegpu.create_nd_tdesc %arg0 : memref<4096x4096xf16> -> !xegpu.tensor_desc<256x32xf16>
diff --git a/mlir/test/python/dialects/transform_xegpu_ext.py b/mlir/test/python/dialects/transform_xegpu_ext.py
index 4f89982ad1c44..2b11acb04ed5b 100644
--- a/mlir/test/python/dialects/transform_xegpu_ext.py
+++ b/mlir/test/python/dialects/transform_xegpu_ext.py
@@ -66,6 +66,25 @@ def setDescLayoutInstData():
     # CHECK: inst_data = [8, 16]
 
 
+@run
+def setDescLayoutSlice():
+    sequence = transform.SequenceOp(
+        transform.FailurePropagationMode.Propagate,
+        [],
+        transform.OperationType.get("xegpu.create_nd_tdesc"),
+    )
+    with InsertionPoint(sequence.body):
+        xegpu.set_desc_layout(
+            sequence.bodyTarget, sg_layout=[6, 4], sg_data=[32, 16], slice_dims=[0]
+        )
+        transform.YieldOp()
+    # CHECK-LABEL: TEST: setDescLayoutSlice
+    # CHECK: %0 = transform.xegpu.set_desc_layout %
+    # CHECK: sg_layout = [6, 4]
+    # CHECK: sg_data = [32, 16]
+    # CHECK: slice_dims = [0]
+
+
 @run
 def setOpLayoutAttrOperandMinimal():
     sequence = transform.SequenceOp(
@@ -106,13 +125,41 @@ def setOpLayoutAttrResult():
             result=True,
         )
         transform.YieldOp()
-    # CHECK-LABEL: TEST: setOpLayoutAttr
+    # CHECK-LABEL: TEST: setOpLayoutAttrResult
+    # CHECK: transform.xegpu.set_op_layout_attr %
+    # NO-CHECK: index = 0
+    # CHECK: result
+    # CHECK: sg_layout = [6, 4]
+    # CHECK: sg_data = [32, 16]
+    # CHECK: inst_data = [8, 16]
+
+
+@run
+def setOpLayoutAttrResultSlice():
+    sequence = transform.SequenceOp(
+        transform.FailurePropagationMode.Propagate,
+        [],
+        transform.OperationType.get("xegpu.dpas"),
+    )
+    with InsertionPoint(sequence.body):
+        xegpu.set_op_layout_attr(
+            sequence.bodyTarget,
+            index=0,
+            sg_layout=[6, 4],
+            sg_data=[32, 16],
+            inst_data=[8, 16],
+            slice_dims=[0],
+            result=True,
+        )
+        transform.YieldOp()
+    # CHECK-LABEL: TEST: setOpLayoutAttrResultSlice
     # CHECK: transform.xegpu.set_op_layout_attr %
     # NO-CHECK: index = 0
     # CHECK: result
     # CHECK: sg_layout = [6, 4]
     # CHECK: sg_data = [32, 16]
     # CHECK: inst_data = [8, 16]
+    # CHECK: slice_dims = [0]
 
 
 @run

@github-actions
Copy link

github-actions bot commented Nov 20, 2025

🐧 Linux x64 Test Results

  • 7103 tests passed
  • 594 tests skipped

Jianhui-Li
Jianhui-Li approved these changes Nov 20, 2025
View reviewed changes
Copy link
Contributor

@Jianhui-Li Jianhui-Li left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM.

@tkarna tkarna merged commit e9fc393 into llvm:main Nov 21, 2025
14 checks passed
@tkarna tkarna deleted the xegpu-tr-ops-slice-layout branch November 21, 2025 08:08
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@rolfmorel rolfmorel rolfmorel approved these changes

@Jianhui-Li Jianhui-Li Jianhui-Li approved these changes

@ftynse ftynse Awaiting requested review from ftynse ftynse is a code owner

@makslevental makslevental Awaiting requested review from makslevental makslevental is a code owner

@stellaraccident stellaraccident Awaiting requested review from stellaraccident stellaraccident is a code owner

@nicolasvasilache nicolasvasilache Awaiting requested review from nicolasvasilache nicolasvasilache is a code owner

@adam-smnk adam-smnk Awaiting requested review from adam-smnk

Assignees

No one assigned

Labels

mlir:gpu mlir:python MLIR Python bindings mlir

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

4 participants

@tkarna @llvmbot @rolfmorel @Jianhui-Li